1 files changed, 116 insertions, 0 deletions
diff --git a/sumdb/internal/tkv/tkvtest/mem.go b/sumdb/internal/tkv/tkvtest/mem.go
new file mode 100644
index 0000000..780edc4
--- /dev/null
+++ b/sumdb/internal/tkv/tkvtest/mem.go
@@ -0,0 +1,116 @@
+// 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
+}