1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
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
}
|
