sumdb/internal/tkv/tkvtest/mem.go

// Copyright 2019 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package tkvtest

import (
	"context"
	"errors"
	"math/rand"
	"sync"

	"golang.org/x/exp/sumdb/internal/tkv"
)

// Mem is an in-memory implementation of Storage.
// It is meant for tests and does not store any data to persistent storage.
//
// The zero value is an empty Mem ready for use.
type Mem struct {
	mu    sync.RWMutex
	table map[string]string
}

// A memTx is a transaction in a Mem.
type memTx struct {
	m      *Mem
	writes []tkv.Write
}

// errRetry is an internal sentinel indicating that the transaction should be retried.
// It is never returned to the caller.
var errRetry = errors.New("retry")

// ReadOnly runs f in a read-only transaction.
func (m *Mem) ReadOnly(ctx context.Context, f func(context.Context, tkv.Transaction) error) error {
	tx := &memTx{m: m}
	for {
		err := func() error {
			m.mu.Lock()
			defer m.mu.Unlock()

			if err := f(ctx, tx); err != nil {
				return err
			}
			// Spurious retry with 10% probability.
			if rand.Intn(10) == 0 {
				return errRetry
			}
			return nil
		}()
		if err != errRetry {
			return err
		}
	}
}

// ReadWrite runs f in a read-write transaction.
func (m *Mem) ReadWrite(ctx context.Context, f func(context.Context, tkv.Transaction) error) error {
	tx := &memTx{m: m}
	for {
		err := func() error {
			m.mu.Lock()
			defer m.mu.Unlock()

			tx.writes = []tkv.Write{}
			if err := f(ctx, tx); err != nil {
				return err
			}
			// Spurious retry with 10% probability.
			if rand.Intn(10) == 0 {
				return errRetry
			}
			if m.table == nil {
				m.table = make(map[string]string)
			}
			for _, w := range tx.writes {
				if w.Value == "" {
					delete(m.table, w.Key)
				} else {
					m.table[w.Key] = w.Value
				}
			}
			return nil
		}()
		if err != errRetry {
			return err
		}
	}
}

// ReadValues returns the values associated with the given keys.
func (tx *memTx) ReadValues(ctx context.Context, keys []string) ([]string, error) {
	vals := make([]string, len(keys))
	for i, key := range keys {
		vals[i] = tx.m.table[key]
	}
	return vals, nil
}

// ReadValue returns the value associated with the single key.
func (tx *memTx) ReadValue(ctx context.Context, key string) (string, error) {
	return tx.m.table[key], nil
}

// BufferWrites buffers a list of writes to be applied
// to the table when the transaction commits.
// The changes are not visible to reads within the transaction.
// The map argument is not used after the call returns.
func (tx *memTx) BufferWrites(list []tkv.Write) error {
	if tx.writes == nil {
		panic("BufferWrite on read-only transaction")
	}
	tx.writes = append(tx.writes, list...)
	return nil
}