package concurrency
import (
v3 "github.com/coreos/etcd/clientv3"
"golang.org/x/net/context"
)
// STM is an interface for software transactional memory.
type STM interface { // Get returns the value for a key and inserts the key in the txn's read set.
// If Get fails, it aborts the transaction with an error, never returning.
Get(key string) string
// Put adds a value for a key to the write set.
Put(key, val string, opts ...v3.OpOption)
// Rev returns the revision of a key in the read set.
Rev(key string) int64
// Del deletes a key.
Del(key string)
// commit attempts to apply the txn's changes to the server.
commit() *v3.TxnResponse
reset()
}
// stmError safely passes STM errors through panic to the STM error channel.
type stmError struct{ err error }
// NewSTMRepeatable initiates new repeatable read transaction; reads within
// the same transaction attempt always return the same data.
func NewSTMRepeatable(ctx context.Context, c *v3.Client, apply func(STM) error) (*v3.TxnResponse, error) { s := &stm{client: c, ctx: ctx, getOpts: []v3.OpOption{v3.WithSerializable()}} return runSTM(s, apply)
}
// NewSTMSerializable initiates a new serialized transaction; reads within the
// same transactiona attempt return data from the revision of the first read.
func NewSTMSerializable(ctx context.Context, c *v3.Client, apply func(STM) error) (*v3.TxnResponse, error) { s := &stmSerializable{ stm: stm{client: c, ctx: ctx}, prefetch: make(map[string]*v3.GetResponse),
}
return runSTM(s, apply)
}
// NewSTMReadCommitted initiates a new read committed transaction.
func NewSTMReadCommitted(ctx context.Context, c *v3.Client, apply func(STM) error) (*v3.TxnResponse, error) { s := &stmReadCommitted{stm{client: c, ctx: ctx, getOpts: []v3.OpOption{v3.WithSerializable()}}} return runSTM(s, apply)
}
type stmResponse struct { resp *v3.TxnResponse
err error
}
func runSTM(s STM, apply func(STM) error) (*v3.TxnResponse, error) { outc := make(chan stmResponse, 1)
go func() { defer func() { if r := recover(); r != nil { e, ok := r.(stmError)
if !ok { // client apply panicked
panic(r)
}
outc <- stmResponse{nil, e.err} }
}()
var out stmResponse
for { s.reset()
if out.err = apply(s); out.err != nil { break
}
if out.resp = s.commit(); out.resp != nil { break
}
}
outc <- out
}()
r := <-outc
return r.resp, r.err
}
// stm implements repeatable-read software transactional memory over etcd
type stm struct { client *v3.Client
ctx context.Context
// rset holds read key values and revisions
rset map[string]*v3.GetResponse
// wset holds overwritten keys and their values
wset map[string]stmPut
// getOpts are the opts used for gets
getOpts []v3.OpOption
}
type stmPut struct { val string
op v3.Op
}
func (s *stm) Get(key string) string { if wv, ok := s.wset[key]; ok { return wv.val
}
return respToValue(s.fetch(key))
}
func (s *stm) Put(key, val string, opts ...v3.OpOption) { s.wset[key] = stmPut{val, v3.OpPut(key, val, opts...)}}
func (s *stm) Del(key string) { s.wset[key] = stmPut{"", v3.OpDelete(key)} }
func (s *stm) Rev(key string) int64 { if resp := s.fetch(key); resp != nil && len(resp.Kvs) != 0 { return resp.Kvs[0].ModRevision
}
return 0
}
func (s *stm) commit() *v3.TxnResponse { txnresp, err := s.client.Txn(s.ctx).If(s.cmps()...).Then(s.puts()...).Commit()
if err != nil { panic(stmError{err}) }
if txnresp.Succeeded { return txnresp
}
return nil
}
// cmps guards the txn from updates to read set
func (s *stm) cmps() []v3.Cmp { cmps := make([]v3.Cmp, 0, len(s.rset))
for k, rk := range s.rset { cmps = append(cmps, isKeyCurrent(k, rk))
}
return cmps
}
func (s *stm) fetch(key string) *v3.GetResponse { if resp, ok := s.rset[key]; ok { return resp
}
resp, err := s.client.Get(s.ctx, key, s.getOpts...)
if err != nil { panic(stmError{err}) }
s.rset[key] = resp
return resp
}
// puts is the list of ops for all pending writes
func (s *stm) puts() []v3.Op { puts := make([]v3.Op, 0, len(s.wset))
for _, v := range s.wset { puts = append(puts, v.op)
}
return puts
}
func (s *stm) reset() { s.rset = make(map[string]*v3.GetResponse)
s.wset = make(map[string]stmPut)
}
type stmSerializable struct { stm
prefetch map[string]*v3.GetResponse
}
func (s *stmSerializable) Get(key string) string { if wv, ok := s.wset[key]; ok { return wv.val
}
firstRead := len(s.rset) == 0
if resp, ok := s.prefetch[key]; ok { delete(s.prefetch, key)
s.rset[key] = resp
}
resp := s.stm.fetch(key)
if firstRead { // txn's base revision is defined by the first read
s.getOpts = []v3.OpOption{ v3.WithRev(resp.Header.Revision),
v3.WithSerializable(),
}
}
return respToValue(resp)
}
func (s *stmSerializable) Rev(key string) int64 { s.Get(key)
return s.stm.Rev(key)
}
func (s *stmSerializable) gets() ([]string, []v3.Op) { keys := make([]string, 0, len(s.rset))
ops := make([]v3.Op, 0, len(s.rset))
for k := range s.rset { keys = append(keys, k)
ops = append(ops, v3.OpGet(k))
}
return keys, ops
}
func (s *stmSerializable) commit() *v3.TxnResponse { keys, getops := s.gets()
txn := s.client.Txn(s.ctx).If(s.cmps()...).Then(s.puts()...)
// use Else to prefetch keys in case of conflict to save a round trip
txnresp, err := txn.Else(getops...).Commit()
if err != nil { panic(stmError{err}) }
if txnresp.Succeeded { return txnresp
}
// load prefetch with Else data
for i := range keys { resp := txnresp.Responses[i].GetResponseRange()
s.rset[keys[i]] = (*v3.GetResponse)(resp)
}
s.prefetch = s.rset
s.getOpts = nil
return nil
}
type stmReadCommitted struct{ stm }
// commit always goes through when read committed
func (s *stmReadCommitted) commit() *v3.TxnResponse { s.rset = nil
return s.stm.commit()
}
func isKeyCurrent(k string, r *v3.GetResponse) v3.Cmp { rev := r.Header.Revision + 1
if len(r.Kvs) != 0 { rev = r.Kvs[0].ModRevision + 1
}
return v3.Compare(v3.ModRevision(k), "<", rev)
}
func respToValue(resp *v3.GetResponse) string { if len(resp.Kvs) == 0 { return ""
}
return string(resp.Kvs[0].Value)
}