sumdb/storage: basic storage interface and tests
Copied from golang.org/x/exp/sumdb/internal/tkv and tkvtest.
For golang/go#31761.
Change-Id: Ib5a411c9b9641d1006a2e721043e8e16bf0f6bea
Reviewed-on: https://go-review.googlesource.com/c/mod/+/176465
Reviewed-by: Hyang-Ah Hana Kim <hyangah@gmail.com>
diff --git a/sumdb/storage/mem.go b/sumdb/storage/mem.go
new file mode 100644
index 0000000..9b7823d
--- /dev/null
+++ b/sumdb/storage/mem.go
@@ -0,0 +1,114 @@
+// 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 storage
+
+import (
+ "context"
+ "errors"
+ "math/rand"
+ "sync"
+)
+
+// 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 []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, 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, Transaction) error) error {
+ tx := &memTx{m: m}
+ for {
+ err := func() error {
+ m.mu.Lock()
+ defer m.mu.Unlock()
+
+ tx.writes = []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 []Write) error {
+ if tx.writes == nil {
+ panic("BufferWrite on read-only transaction")
+ }
+ tx.writes = append(tx.writes, list...)
+ return nil
+}
diff --git a/sumdb/storage/mem_test.go b/sumdb/storage/mem_test.go
new file mode 100644
index 0000000..57ae15e
--- /dev/null
+++ b/sumdb/storage/mem_test.go
@@ -0,0 +1,14 @@
+// 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 storage
+
+import (
+ "context"
+ "testing"
+)
+
+func TestMem(t *testing.T) {
+ TestStorage(t, context.Background(), new(Mem))
+}
diff --git a/sumdb/storage/storage.go b/sumdb/storage/storage.go
new file mode 100644
index 0000000..3c80ac2
--- /dev/null
+++ b/sumdb/storage/storage.go
@@ -0,0 +1,56 @@
+// 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 storage defines storage interfaces for and a basic implementation of a checksum database.
+package storage
+
+import "context"
+
+// A Storage is a transaction key-value storage system.
+type Storage interface {
+ // ReadOnly runs f in a read-only transaction.
+ // It is equivalent to ReadWrite except that the
+ // transaction's BufferWrite method will fail unconditionally.
+ // (The implementation may be able to optimize the
+ // transaction if it knows at the start that no writes will happen.)
+ ReadOnly(ctx context.Context, f func(context.Context, Transaction) error) error
+
+ // ReadWrite runs f in a read-write transaction.
+ // If f returns an error, the transaction aborts and returns that error.
+ // If f returns nil, the transaction attempts to commit and then then return nil.
+ // Otherwise it tries again. Note that f may be called multiple times and that
+ // the result only describes the effect of the final call to f.
+ // The caller must take care not to use any state computed during
+ // earlier calls to f, or even the last call to f when an error is returned.
+ ReadWrite(ctx context.Context, f func(context.Context, Transaction) error) error
+}
+
+// A Transaction provides read and write operations within a transaction,
+// as executed by Storage's ReadOnly or ReadWrite methods.
+type Transaction interface {
+ // ReadValue reads the value associated with a single key.
+ // If there is no value associated with that key, ReadKey returns an empty value.
+ // An error is only returned for problems accessing the storage.
+ ReadValue(ctx context.Context, key string) (value string, err error)
+
+ // ReadValues reads the values associated with the given keys.
+ // If there is no value stored for a given key, ReadValues returns an empty value for that key.
+ // An error is only returned for problems accessing the storage.
+ ReadValues(ctx context.Context, keys []string) (values []string, err error)
+
+ // BufferWrites buffers the given writes,
+ // to be applied at the end of the transaction.
+ // BufferWrites panics if this is a ReadOnly transaction.
+ // It returns an error if it detects any other problems.
+ // The behavior of multiple writes buffered using the same key
+ // is undefined: it may return an error or not.
+ BufferWrites(writes []Write) error
+}
+
+// A Write is a single change to be applied at the end of a read-write transaction.
+// A Write with an empty value deletes the value associated with the given key.
+type Write struct {
+ Key string
+ Value string
+}
diff --git a/sumdb/storage/test.go b/sumdb/storage/test.go
new file mode 100644
index 0000000..fdf410e
--- /dev/null
+++ b/sumdb/storage/test.go
@@ -0,0 +1,75 @@
+// 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 storage
+
+import (
+ "context"
+ "fmt"
+ "io"
+ "testing"
+)
+
+// TestStorage tests a Storage implementation.
+func TestStorage(t *testing.T, ctx context.Context, storage Storage) {
+ s := storage
+
+ // Insert records.
+ err := s.ReadWrite(ctx, func(ctx context.Context, tx Transaction) error {
+ for i := 0; i < 10; i++ {
+ err := tx.BufferWrites([]Write{
+ {Key: fmt.Sprint(i), Value: fmt.Sprint(-i)},
+ {Key: fmt.Sprint(1000 + i), Value: fmt.Sprint(-1000 - i)},
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+ }
+ return nil
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+
+ // Read the records back.
+ testRead := func() {
+ err := s.ReadOnly(ctx, func(ctx context.Context, tx Transaction) error {
+ for i := int64(0); i < 1010; i++ {
+ if i == 10 {
+ i = 1000
+ }
+ val, err := tx.ReadValue(ctx, fmt.Sprint(i))
+ if err != nil {
+ t.Fatalf("reading %v: %v", i, err)
+ }
+ if want := fmt.Sprint(-i); val != want {
+ t.Fatalf("ReadValue %v = %q, want %v", i, val, want)
+ }
+ }
+ return nil
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+ }
+ testRead()
+
+ // Buffered writes in failed transaction should not be applied.
+ err = s.ReadWrite(ctx, func(ctx context.Context, tx Transaction) error {
+ tx.BufferWrites([]Write{
+ {Key: fmt.Sprint(0), Value: ""}, // delete
+ {Key: fmt.Sprint(1), Value: "overwrite"}, // overwrite
+ })
+ if err != nil {
+ t.Fatal(err)
+ }
+ return io.ErrUnexpectedEOF
+ })
+ if err != io.ErrUnexpectedEOF {
+ t.Fatalf("ReadWrite returned %v, want ErrUnexpectedEOF", err)
+ }
+
+ // All same values should still be there.
+ testRead()
+}
