consultopo: make retry backoff context-aware

Replace time.Sleep with a select on ctx.Done and time.After so that
retries abort promptly when the caller's context is canceled or
expired. Pass context from callers in file.go, directory.go, and
election.go via QueryOptions so the retry loop can observe it.

Co-Authored-By: Claude <svc-devxp-claude@slack-corp.com>
Signed-off-by: Steve Baker <s.baker@slack-corp.com>

Author: sbaker617

go/vt/topo/consultopo/directory.go

@@ -17,10 +17,11 @@ limitations under the License.
 package consultopo
 
 import (
+	"context"
 	"path"
 	"strings"
 
-	"context"
+	"github.com/hashicorp/consul/api"
 
 	"vitess.io/vitess/go/vt/topo"
 )
@@ -36,7 +37,7 @@ func (s *Server) ListDir(ctx context.Context, dirPath string, full bool) ([]topo
 
 	isRoot := dirPath == "" || dirPath == "/"
 
-	keys, _, err := s.kv.Keys(nodePath, "", nil)
+	keys, _, err := s.kv.Keys(nodePath, "", (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		return nil, err
 	}

go/vt/topo/consultopo/election.go

@@ -123,7 +123,7 @@ func (mp *consulLeaderParticipation) Stop() {
 // GetCurrentLeaderID is part of the topo.LeaderParticipation interface
 func (mp *consulLeaderParticipation) GetCurrentLeaderID(ctx context.Context) (string, error) {
 	electionPath := path.Join(mp.s.root, electionsPath, mp.name)
-	pair, _, err := mp.s.kv.Get(electionPath, nil)
+	pair, _, err := mp.s.kv.Get(electionPath, (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		return "", err
 	}

go/vt/topo/consultopo/file.go

@@ -17,9 +17,8 @@ limitations under the License.
 package consultopo
 
 import (
-	"path"
-
 	"context"
+	"path"
 
 	"github.com/hashicorp/consul/api"
 
@@ -41,7 +40,7 @@ func (s *Server) Create(ctx context.Context, filePath string, contents []byte) (
 			Index: 0,
 		},
 	}
-	ok, resp, _, err := s.kv.Txn(ops, nil)
+	ok, resp, _, err := s.kv.Txn(ops, (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		// Communication error.
 		return nil, err
@@ -70,7 +69,7 @@ func (s *Server) Update(ctx context.Context, filePath string, contents []byte, v
 		ops[0].Verb = api.KVCAS
 		ops[0].Index = uint64(version.(ConsulVersion))
 	}
-	ok, resp, _, err := s.kv.Txn(ops, nil)
+	ok, resp, _, err := s.kv.Txn(ops, (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		// Communication error.
 		return nil, err
@@ -87,7 +86,7 @@ func (s *Server) Update(ctx context.Context, filePath string, contents []byte, v
 func (s *Server) Get(ctx context.Context, filePath string) ([]byte, topo.Version, error) {
 	nodePath := path.Join(s.root, filePath)
 
-	pair, _, err := s.kv.Get(nodePath, nil)
+	pair, _, err := s.kv.Get(nodePath, (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		return nil, nil, err
 	}
@@ -107,7 +106,7 @@ func (s *Server) GetVersion(ctx context.Context, filePath string, version int64)
 func (s *Server) List(ctx context.Context, filePathPrefix string) ([]topo.KVInfo, error) {
 	nodePathPrefix := path.Join(s.root, filePathPrefix)
 
-	pairs, _, err := s.kv.List(nodePathPrefix, nil)
+	pairs, _, err := s.kv.List(nodePathPrefix, (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		return []topo.KVInfo{}, err
 	}
@@ -150,7 +149,7 @@ func (s *Server) Delete(ctx context.Context, filePath string, version topo.Versi
 		ops[1].Verb = api.KVDeleteCAS
 		ops[1].Index = uint64(version.(ConsulVersion))
 	}
-	ok, resp, _, err := s.kv.Txn(ops, nil)
+	ok, resp, _, err := s.kv.Txn(ops, (&api.QueryOptions{}).WithContext(ctx))
 	if err != nil {
 		// Communication error.
 		return err

go/vt/topo/consultopo/retry.go

@@ -68,7 +68,8 @@ func newRetryKV(inner kvClient, count int, baseDelay, maxDelay time.Duration, en
 func (r *retryKV) Get(key string, q *api.QueryOptions) (*api.KVPair, *api.QueryMeta, error) {
 	var pair *api.KVPair
 	var meta *api.QueryMeta
-	err := r.retry(func() error {
+	ctx := contextFromQueryOptions(q)
+	err := r.retry(ctx, func() error {
 		var ierr error
 		pair, meta, ierr = r.inner.Get(key, q)
 		return ierr
@@ -79,7 +80,8 @@ func (r *retryKV) Get(key string, q *api.QueryOptions) (*api.KVPair, *api.QueryM
 func (r *retryKV) List(prefix string, q *api.QueryOptions) (api.KVPairs, *api.QueryMeta, error) {
 	var pairs api.KVPairs
 	var meta *api.QueryMeta
-	err := r.retry(func() error {
+	ctx := contextFromQueryOptions(q)
+	err := r.retry(ctx, func() error {
 		var ierr error
 		pairs, meta, ierr = r.inner.List(prefix, q)
 		return ierr
@@ -90,7 +92,8 @@ func (r *retryKV) List(prefix string, q *api.QueryOptions) (api.KVPairs, *api.Qu
 func (r *retryKV) Keys(prefix string, separator string, q *api.QueryOptions) ([]string, *api.QueryMeta, error) {
 	var keys []string
 	var meta *api.QueryMeta
-	err := r.retry(func() error {
+	ctx := contextFromQueryOptions(q)
+	err := r.retry(ctx, func() error {
 		var ierr error
 		keys, meta, ierr = r.inner.Keys(prefix, separator, q)
 		return ierr
@@ -102,15 +105,23 @@ func (r *retryKV) Txn(txn api.KVTxnOps, q *api.QueryOptions) (bool, *api.KVTxnRe
 	var ok bool
 	var resp *api.KVTxnResponse
 	var meta *api.QueryMeta
-	err := r.retry(func() error {
+	ctx := contextFromQueryOptions(q)
+	err := r.retry(ctx, func() error {
 		var ierr error
 		ok, resp, meta, ierr = r.inner.Txn(txn, q)
 		return ierr
 	})
 	return ok, resp, meta, err
 }
 
-func (r *retryKV) retry(action func() error) error {
+func contextFromQueryOptions(q *api.QueryOptions) context.Context {
+	if q != nil && q.Context() != nil {
+		return q.Context()
+	}
+	return context.Background()
+}
+
+func (r *retryKV) retry(ctx context.Context, action func() error) error {
 	if !r.enabled {
 		return action()
 	}
@@ -119,7 +130,11 @@ func (r *retryKV) retry(action func() error) error {
 	for attempt := 0; attempt < r.count; attempt++ {
 		if attempt > 0 {
 			delay := r.backoff(attempt)
-			time.Sleep(delay)
+			select {
+			case <-ctx.Done():
+				return ctx.Err()
+			case <-time.After(delay):
+			}
 		}
 
 		err = action()

go/vt/topo/consultopo/retry_test.go

@@ -246,6 +246,28 @@ func TestRetryKV_Backoff(t *testing.T) {
 	assert.LessOrEqual(t, d10, r.maxDelay+r.maxDelay/4)
 }
 
+func TestRetryKV_Get_ContextCanceledDuringBackoff(t *testing.T) {
+	ctx, cancel := context.WithCancel(context.Background())
+	mock := &mockKV{
+		getFunc: func(call int) (*api.KVPair, *api.QueryMeta, error) {
+			if call == 1 {
+				cancel()
+			}
+			return nil, nil, errors.New("Unexpected response code: 500")
+		},
+	}
+	r := newRetryKV(mock, 3, 500*time.Millisecond, 5*time.Second, true)
+
+	opts := (&api.QueryOptions{}).WithContext(ctx)
+	start := time.Now()
+	_, _, err := r.Get("test", opts)
+	elapsed := time.Since(start)
+
+	assert.ErrorIs(t, err, context.Canceled)
+	assert.Equal(t, 1, mock.getCalls)
+	assert.Less(t, elapsed, 200*time.Millisecond)
+}
+
 func TestRetryKV_Backoff_ZeroBaseDelay(t *testing.T) {
 	r := &retryKV{
 		baseDelay: 0,