Browse Source

concurrency: STM WithPrefetch option

Fixes #6923
Anthony Romano 9 years ago
parent
commit
8604d1863b
1 changed files with 18 additions and 2 deletions
  1. 18 2
      clientv3/concurrency/stm.go

+ 18 - 2
clientv3/concurrency/stm.go

@@ -57,8 +57,9 @@ const (
 type stmError struct{ err error }
 
 type stmOptions struct {
-	iso Isolation
-	ctx context.Context
+	iso      Isolation
+	ctx      context.Context
+	prefetch []string
 }
 
 type stmOption func(*stmOptions)
@@ -73,12 +74,27 @@ func WithAbortContext(ctx context.Context) stmOption {
 	return func(so *stmOptions) { so.ctx = ctx }
 }
 
+// WithPrefetch is a hint to prefetch a list of keys before trying to apply.
+// If an STM transaction will unconditionally fetch a set of keys, prefetching
+// those keys will save the round-trip cost from requesting each key one by one
+// with Get().
+func WithPrefetch(keys ...string) stmOption {
+	return func(so *stmOptions) { so.prefetch = append(so.prefetch, keys...) }
+}
+
 // NewSTM initiates a new STM instance.
 func NewSTM(c *v3.Client, apply func(STM) error, so ...stmOption) (*v3.TxnResponse, error) {
 	opts := &stmOptions{ctx: c.Ctx()}
 	for _, f := range so {
 		f(opts)
 	}
+	if len(opts.prefetch) != 0 {
+		f := apply
+		apply = func(s STM) error {
+			s.Get(opts.prefetch...)
+			return f(s)
+		}
+	}
 	var s STM
 	switch opts.iso {
 	case Serializable: