refactor(kv): plumb caller context through write + verify-leader paths

adapter/distribution_server.go

@@ -139,7 +139,7 @@ func (s *DistributionServer) SplitRange(ctx context.Context, req *pb.SplitRangeR
 	s.mu.Lock()
 	defer s.mu.Unlock()
 
-	if err := s.verifyCatalogLeader(); err != nil {
+	if err := s.verifyCatalogLeader(ctx); err != nil {
 		return nil, err
 	}
 
@@ -191,15 +191,15 @@ func (s *DistributionServer) pinReadTS(ts uint64) *kv.ActiveTimestampToken {
 	return s.readTracker.Pin(ts)
 }
 
-func (s *DistributionServer) verifyCatalogLeader() error {
+func (s *DistributionServer) verifyCatalogLeader(ctx context.Context) error {
 	if s.coordinator == nil {
 		return grpcStatusError(codes.FailedPrecondition, errDistributionCoordinatorRequired.Error())
 	}
 	key := distribution.CatalogVersionKey()
 	if !s.coordinator.IsLeaderForKey(key) {
 		return grpcStatusError(codes.FailedPrecondition, errDistributionNotLeader.Error())
 	}
-	if err := s.coordinator.VerifyLeaderForKey(key); err != nil {
+	if err := s.coordinator.VerifyLeaderForKey(ctx, key); err != nil {
 		return grpcStatusErrorf(codes.FailedPrecondition, "verify catalog leader: %v", err)
 	}
 	return nil

adapter/distribution_server_test.go

@@ -713,7 +713,7 @@ func (s *distributionCoordinatorStub) IsLeader() bool {
 	return s.leader
 }
 
-func (s *distributionCoordinatorStub) VerifyLeader() error {
+func (s *distributionCoordinatorStub) VerifyLeader(context.Context) error {
 	if !s.leader {
 		return kv.ErrLeaderNotFound
 	}
@@ -728,7 +728,7 @@ func (s *distributionCoordinatorStub) IsLeaderForKey(_ []byte) bool {
 	return s.leader
 }
 
-func (s *distributionCoordinatorStub) VerifyLeaderForKey(_ []byte) error {
+func (s *distributionCoordinatorStub) VerifyLeaderForKey(_ context.Context, _ []byte) error {
 	if !s.leader {
 		return kv.ErrLeaderNotFound
 	}

adapter/dynamodb.go

@@ -377,19 +377,19 @@ func (d *DynamoDBServer) serveDynamoLeaderHealthz(w http.ResponseWriter, r *http
 		return
 	}
 
-	if isVerifiedDynamoLeader(d.coordinator) {
+	if isVerifiedDynamoLeader(r.Context(), d.coordinator) {
 		writeDynamoHealthBody(w, r, http.StatusOK, "ok\n")
 		return
 	}
 
 	writeDynamoHealthBody(w, r, http.StatusServiceUnavailable, "not leader\n")
 }
 
-func isVerifiedDynamoLeader(coordinator kv.Coordinator) bool {
+func isVerifiedDynamoLeader(ctx context.Context, coordinator kv.Coordinator) bool {
 	if coordinator == nil || !coordinator.IsLeader() {
 		return false
 	}
-	return coordinator.VerifyLeader() == nil
+	return coordinator.VerifyLeader(ctx) == nil
 }
 
 func writeDynamoHealthMethod(w http.ResponseWriter, r *http.Request) bool {

adapter/dynamodb_admin.go

@@ -202,7 +202,7 @@ func (d *DynamoDBServer) AdminCreateTable(ctx context.Context, principal AdminPr
 	if !principal.Role.canWrite() {
 		return nil, ErrAdminForbidden
 	}
-	if !isVerifiedDynamoLeader(d.coordinator) {
+	if !isVerifiedDynamoLeader(ctx, d.coordinator) {
 		return nil, ErrAdminNotLeader
 	}
 	legacy, err := buildLegacyCreateTableInput(in)
@@ -249,7 +249,7 @@ func (d *DynamoDBServer) AdminDeleteTable(ctx context.Context, principal AdminPr
 	if !principal.Role.canWrite() {
 		return ErrAdminForbidden
 	}
-	if !isVerifiedDynamoLeader(d.coordinator) {
+	if !isVerifiedDynamoLeader(ctx, d.coordinator) {
 		return ErrAdminNotLeader
 	}
 	if strings.TrimSpace(name) == "" {

adapter/dynamodb_test.go

@@ -1831,8 +1831,8 @@ func (w *testCoordinatorWrapper) IsLeader() bool {
 	return w.inner.IsLeader()
 }
 
-func (w *testCoordinatorWrapper) VerifyLeader() error {
-	return w.inner.VerifyLeader()
+func (w *testCoordinatorWrapper) VerifyLeader(ctx context.Context) error {
+	return w.inner.VerifyLeader(ctx)
 }
 
 func (w *testCoordinatorWrapper) RaftLeader() string {
@@ -1843,8 +1843,8 @@ func (w *testCoordinatorWrapper) IsLeaderForKey(key []byte) bool {
 	return w.inner.IsLeaderForKey(key)
 }
 
-func (w *testCoordinatorWrapper) VerifyLeaderForKey(key []byte) error {
-	return w.inner.VerifyLeaderForKey(key)
+func (w *testCoordinatorWrapper) VerifyLeaderForKey(ctx context.Context, key []byte) error {
+	return w.inner.VerifyLeaderForKey(ctx, key)
 }
 
 func (w *testCoordinatorWrapper) RaftLeaderForKey(key []byte) string {

adapter/internal.go

@@ -49,7 +49,7 @@ func (i *Internal) Forward(ctx context.Context, req *pb.ForwardRequest) (*pb.For
 		}, errors.WithStack(err)
 	}
 
-	r, err := i.transactionManager.Commit(req.Requests)
+	r, err := i.transactionManager.Commit(ctx, req.Requests)
 	if err != nil {
 		return &pb.ForwardResponse{
 			Success:     false,

adapter/redis.go

@@ -1287,7 +1287,14 @@ func (r *RedisServer) keys(conn redcon.Conn, cmd redcon.Command) {
 	pattern := cmd.Args[1]
 
 	if r.coordinator.IsLeader() {
-		if err := r.coordinator.VerifyLeader(); err != nil {
+		// Per-call ctx with redisDispatchTimeout instead of the
+		// long-lived handlerContext: a stalled VerifyLeader on KEYS
+		// must not pin the command handler indefinitely. The same
+		// bound the rest of the dispatch path (sadd, set, …) uses;
+		// see Codex P1 review on PR #749.
+		ctx, cancel := context.WithTimeout(r.handlerContext(), redisDispatchTimeout)
+		defer cancel()
+		if err := r.coordinator.VerifyLeader(ctx); err != nil {
 			conn.WriteError(err.Error())
 			return
 		}
@@ -3254,7 +3261,7 @@ func (r *RedisServer) rangeList(key []byte, startRaw, endRaw []byte) ([]string,
 		return nil, wrongTypeError()
 	}
 
-	if err := r.coordinator.VerifyLeaderForKey(key); err != nil {
+	if err := r.coordinator.VerifyLeaderForKey(r.handlerContext(), key); err != nil {
 		return nil, errors.WithStack(err)
 	}
 
@@ -3510,7 +3517,7 @@ func (r *RedisServer) readValueAt(key []byte, readTS uint64) ([]byte, error) {
 	}
 
 	if r.coordinator.IsLeaderForKey(key) {
-		if err := r.coordinator.VerifyLeaderForKey(key); err != nil {
+		if err := r.coordinator.VerifyLeaderForKey(r.handlerContext(), key); err != nil {
 			return nil, errors.WithStack(err)
 		}
 		v, err := r.store.GetAt(context.Background(), key, readTS)

adapter/redis_compat_commands.go

@@ -1039,7 +1039,7 @@ func (r *RedisServer) dbsize(conn redcon.Conn, _ redcon.Command) {
 		conn.WriteInt(size)
 		return
 	}
-	if err := r.coordinator.VerifyLeader(); err != nil {
+	if err := r.coordinator.VerifyLeader(r.handlerContext()); err != nil {
 		conn.WriteError(err.Error())
 		return
 	}
@@ -1144,7 +1144,11 @@ func (r *RedisServer) flushDatabase(conn redcon.Conn, all bool) {
 	defer cancel()
 
 	if err := r.retryRedisWrite(ctx, func() error {
-		if err := r.coordinator.VerifyLeader(); err != nil {
+		// Use the per-call ctx with redisDispatchTimeout, NOT
+		// handlerContext (the long-lived server baseCtx). FLUSHDB's
+		// retry budget already lives in ctx; routing it to
+		// VerifyLeader keeps the whole command bounded.
+		if err := r.coordinator.VerifyLeader(ctx); err != nil {
 			return fmt.Errorf("verify leader: %w", err)
 		}
 

adapter/redis_compat_helpers.go

@@ -860,7 +860,7 @@ func (r *RedisServer) doGetAt(key []byte, readTS uint64, verify bool) ([]byte, e
 	}
 	if r.coordinator.IsLeaderForKey(routingKey) {
 		if verify {
-			if err := r.coordinator.VerifyLeaderForKey(routingKey); err != nil {
+			if err := r.coordinator.VerifyLeaderForKey(r.handlerContext(), routingKey); err != nil {
 				return nil, errors.WithStack(err)
 			}
 		}

adapter/redis_hello_test.go

@@ -98,12 +98,12 @@ type helloTestCoordinator struct {
 func (c *helloTestCoordinator) Dispatch(context.Context, *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) {
 	return &kv.CoordinateResponse{}, nil
 }
-func (c *helloTestCoordinator) IsLeader() bool                  { return c.isLeader }
-func (c *helloTestCoordinator) VerifyLeader() error             { return nil }
-func (c *helloTestCoordinator) RaftLeader() string              { return "" }
-func (c *helloTestCoordinator) IsLeaderForKey([]byte) bool      { return c.isLeader }
-func (c *helloTestCoordinator) VerifyLeaderForKey([]byte) error { return nil }
-func (c *helloTestCoordinator) RaftLeaderForKey([]byte) string  { return "" }
+func (c *helloTestCoordinator) IsLeader() bool                                   { return c.isLeader }
+func (c *helloTestCoordinator) VerifyLeader(context.Context) error               { return nil }
+func (c *helloTestCoordinator) RaftLeader() string                               { return "" }
+func (c *helloTestCoordinator) IsLeaderForKey([]byte) bool                       { return c.isLeader }
+func (c *helloTestCoordinator) VerifyLeaderForKey(context.Context, []byte) error { return nil }
+func (c *helloTestCoordinator) RaftLeaderForKey([]byte) string                   { return "" }
 func (c *helloTestCoordinator) Clock() *kv.HLC {
 	if c.clock == nil {
 		c.clock = kv.NewHLC()

adapter/redis_info_test.go

@@ -20,12 +20,12 @@ type infoTestCoordinator struct {
 func (c *infoTestCoordinator) Dispatch(context.Context, *kv.OperationGroup[kv.OP]) (*kv.CoordinateResponse, error) {
 	return &kv.CoordinateResponse{}, nil
 }
-func (c *infoTestCoordinator) IsLeader() bool                  { return c.isLeader }
-func (c *infoTestCoordinator) VerifyLeader() error             { return nil }
-func (c *infoTestCoordinator) RaftLeader() string              { return c.raftLeader }
-func (c *infoTestCoordinator) IsLeaderForKey([]byte) bool      { return c.isLeader }
-func (c *infoTestCoordinator) VerifyLeaderForKey([]byte) error { return nil }
-func (c *infoTestCoordinator) RaftLeaderForKey([]byte) string  { return c.raftLeader }
+func (c *infoTestCoordinator) IsLeader() bool                                   { return c.isLeader }
+func (c *infoTestCoordinator) VerifyLeader(context.Context) error               { return nil }
+func (c *infoTestCoordinator) RaftLeader() string                               { return c.raftLeader }
+func (c *infoTestCoordinator) IsLeaderForKey([]byte) bool                       { return c.isLeader }
+func (c *infoTestCoordinator) VerifyLeaderForKey(context.Context, []byte) error { return nil }
+func (c *infoTestCoordinator) RaftLeaderForKey([]byte) string                   { return c.raftLeader }
 func (c *infoTestCoordinator) Clock() *kv.HLC {
 	if c.clock == nil {
 		c.clock = kv.NewHLC()

adapter/redis_keys_pattern_test.go

@@ -29,7 +29,7 @@ func (s *stubAdapterCoordinator) IsLeader() bool {
 	return true
 }
 
-func (s *stubAdapterCoordinator) VerifyLeader() error {
+func (s *stubAdapterCoordinator) VerifyLeader(context.Context) error {
 	s.verifyCalls.Add(1)
 	return s.verifyLeaderErr
 }
@@ -45,7 +45,7 @@ func (s *stubAdapterCoordinator) IsLeaderForKey([]byte) bool {
 	return true
 }
 
-func (s *stubAdapterCoordinator) VerifyLeaderForKey([]byte) error {
+func (s *stubAdapterCoordinator) VerifyLeaderForKey(context.Context, []byte) error {
 	return nil
 }
 

adapter/redis_retry_test.go

@@ -56,7 +56,7 @@ func (c *retryOnceCoordinator) IsLeader() bool {
 	return true
 }
 
-func (c *retryOnceCoordinator) VerifyLeader() error {
+func (c *retryOnceCoordinator) VerifyLeader(context.Context) error {
 	return nil
 }
 
@@ -68,7 +68,7 @@ func (c *retryOnceCoordinator) IsLeaderForKey([]byte) bool {
 	return true
 }
 
-func (c *retryOnceCoordinator) VerifyLeaderForKey([]byte) error {
+func (c *retryOnceCoordinator) VerifyLeaderForKey(context.Context, []byte) error {
 	return nil
 }
 

adapter/s3.go

@@ -2297,12 +2297,12 @@ func (s *S3Server) maybeProxyToLeader(w http.ResponseWriter, r *http.Request) bo
 	}
 	var leader string
 	if len(key) > 0 {
-		if s.coordinator.IsLeaderForKey(key) && s.coordinator.VerifyLeaderForKey(key) == nil {
+		if s.coordinator.IsLeaderForKey(key) && s.coordinator.VerifyLeaderForKey(r.Context(), key) == nil {
 			return false
 		}
 		leader = s.coordinator.RaftLeaderForKey(key)
 	} else {
-		if s.coordinator.IsLeader() && s.coordinator.VerifyLeader() == nil {
+		if s.coordinator.IsLeader() && s.coordinator.VerifyLeader(r.Context()) == nil {
 			return false
 		}
 		leader = s.coordinator.RaftLeader()
@@ -2420,7 +2420,7 @@ func (s *S3Server) serveS3LeaderHealthz(w http.ResponseWriter, r *http.Request)
 		return true
 	}
 	status, body := http.StatusOK, "ok"
-	if !s.isVerifiedS3Leader() {
+	if !s.isVerifiedS3Leader(r.Context()) {
 		status, body = http.StatusServiceUnavailable, "not leader"
 	}
 	w.WriteHeader(status)
@@ -2430,11 +2430,11 @@ func (s *S3Server) serveS3LeaderHealthz(w http.ResponseWriter, r *http.Request)
 	return true
 }
 
-func (s *S3Server) isVerifiedS3Leader() bool {
+func (s *S3Server) isVerifiedS3Leader(ctx context.Context) bool {
 	if s.coordinator == nil || !s.coordinator.IsLeader() {
 		return false
 	}
-	return s.coordinator.VerifyLeader() == nil
+	return s.coordinator.VerifyLeader(ctx) == nil
 }
 
 // prepareStreamingPutBody wraps r.Body for aws-chunked framed uploads. When

adapter/s3_admin.go

@@ -226,7 +226,7 @@ func (s *S3Server) AdminCreateBucket(ctx context.Context, principal AdminPrincip
 	if !principal.Role.canWrite() {
 		return nil, ErrAdminForbidden
 	}
-	if !s.isVerifiedS3Leader() {
+	if !s.isVerifiedS3Leader(ctx) {
 		return nil, ErrAdminNotLeader
 	}
 	if err := validateS3BucketName(name); err != nil {
@@ -314,7 +314,7 @@ func (s *S3Server) AdminPutBucketAcl(ctx context.Context, principal AdminPrincip
 	if !principal.Role.canWrite() {
 		return ErrAdminForbidden
 	}
-	if !s.isVerifiedS3Leader() {
+	if !s.isVerifiedS3Leader(ctx) {
 		return ErrAdminNotLeader
 	}
 	acl = adminCanonicalACL(acl)
@@ -406,7 +406,7 @@ func (s *S3Server) AdminDeleteBucket(ctx context.Context, principal AdminPrincip
 	if !principal.Role.canWrite() {
 		return ErrAdminForbidden
 	}
-	if !s.isVerifiedS3Leader() {
+	if !s.isVerifiedS3Leader(ctx) {
 		return ErrAdminNotLeader
 	}
 

adapter/s3_test.go

@@ -696,7 +696,7 @@ func (c *followerS3Coordinator) IsLeader() bool {
 	return false
 }
 
-func (c *followerS3Coordinator) VerifyLeader() error {
+func (c *followerS3Coordinator) VerifyLeader(context.Context) error {
 	return kv.ErrLeaderNotFound
 }
 
@@ -729,7 +729,7 @@ func (c *routeAwareS3Coordinator) IsLeaderForKey(key []byte) bool {
 	return c.localForKey(key)
 }
 
-func (c *routeAwareS3Coordinator) VerifyLeaderForKey(key []byte) error {
+func (c *routeAwareS3Coordinator) VerifyLeaderForKey(_ context.Context, key []byte) error {
 	if c.IsLeaderForKey(key) {
 		return nil
 	}

adapter/sqs.go

@@ -498,18 +498,18 @@ func (s *SQSServer) serveSQSLeaderHealthz(w http.ResponseWriter, r *http.Request
 	if !writeSQSHealthMethod(w, r) {
 		return
 	}
-	if isVerifiedSQSLeader(s.coordinator) {
+	if isVerifiedSQSLeader(r.Context(), s.coordinator) {
 		writeSQSHealthBody(w, r, http.StatusOK, "ok\n")
 		return
 	}
 	writeSQSHealthBody(w, r, http.StatusServiceUnavailable, "not leader\n")
 }
 
-func isVerifiedSQSLeader(coordinator kv.Coordinator) bool {
+func isVerifiedSQSLeader(ctx context.Context, coordinator kv.Coordinator) bool {
 	if coordinator == nil || !coordinator.IsLeader() {
 		return false
 	}
-	return coordinator.VerifyLeader() == nil
+	return coordinator.VerifyLeader(ctx) == nil
 }
 
 func writeSQSHealthMethod(w http.ResponseWriter, r *http.Request) bool {

adapter/sqs_admin.go

@@ -105,7 +105,7 @@ func (s *SQSServer) AdminDeleteQueue(ctx context.Context, principal AdminPrincip
 	if !principal.Role.canWrite() {
 		return ErrAdminForbidden
 	}
-	if !isVerifiedSQSLeader(s.coordinator) {
+	if !isVerifiedSQSLeader(ctx, s.coordinator) {
 		return ErrAdminNotLeader
 	}
 	if strings.TrimSpace(name) == "" {