From 369362eb55847f64191b7c19ec13569d3d4eb300 Mon Sep 17 00:00:00 2001 From: Ovais Tariq Date: Fri, 25 Sep 2026 20:22:34 -0700 Subject: [PATCH 1/2] fix: dial peers outside the router lock and honour caller context The router dialed a peer while holding its router-wide write lock, with a blocking dial of up to ConnectionTimeout (5s). One unreachable peer, the normal state of a node mid-restart, therefore stalled every remote Put, Get and Delete to every other node behind that dial, for up to four attempts per request until the circuit breaker opened. Dials now run under a per-node lock only: Router.mu guards just the clients map, callers for the same node share one attempt, and callers for other nodes are unaffected. A dial that lands after RemoveClient or RefreshConnections dropped the node is closed rather than leaked. Route also had no context: retries slept with time.Sleep and the dial was bound to context.Background, so a request the client had already abandoned kept dialing and backing off against the owner. RouteContext binds both to the caller's context and the remote paths in server/operations and the streaming forwarders pass theirs. An abandoned dial is not counted against the node's circuit breaker. Route and RouteWithRetry keep their signatures. Perfloop cases case_vzhjv4f9s4 and case_yf8b86gxc0. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01YCoN4Him5ZY4hZovVp37ra Signed-off-by: Ovais Tariq --- coordinator/coordinator.go | 5 + coordinator/router.go | 186 +++++++++++++++++++++++--------- coordinator/router_dial_test.go | 134 +++++++++++++++++++++++ server/operations/cas.go | 10 +- server/operations/delete.go | 2 +- server/operations/get.go | 2 +- server/operations/operations.go | 11 ++ server/operations/put.go | 2 +- server/service/cas.go | 4 +- server/service/get.go | 2 +- server/service/put.go | 2 +- 11 files changed, 296 insertions(+), 64 deletions(-) create mode 100644 coordinator/router_dial_test.go diff --git a/coordinator/coordinator.go b/coordinator/coordinator.go index b3b066db..60a657b3 100644 --- a/coordinator/coordinator.go +++ b/coordinator/coordinator.go @@ -340,6 +340,11 @@ func (c *Coordinator) Route(key string) (pb.CacheServiceClient, error) { return c.router.Route(key) } +// RouteContext is Route bound to the caller's context; see Router.RouteContext. +func (c *Coordinator) RouteContext(ctx context.Context, key string) (pb.CacheServiceClient, error) { + return c.router.RouteContext(ctx, key) +} + // GetNodeForKey returns the node for the given key func (c *Coordinator) GetNodeForKey(key string) (*ring.NodeInfo, error) { return c.ringManager.GetNode(key) diff --git a/coordinator/router.go b/coordinator/router.go index 83fccaed..1c56493a 100644 --- a/coordinator/router.go +++ b/coordinator/router.go @@ -79,13 +79,28 @@ func DefaultRouterConfig() *RouterConfig { // clientState tracks the state of a client connection type clientState struct { - client pb.CacheServiceClient - conn *grpc.ClientConn + client pb.CacheServiceClient // guarded by mu + conn *grpc.ClientConn // guarded by mu failureCount int32 circuitOpenTime time.Time circuitOpen int32 // atomic: 0=closed, 1=open lastFailure time.Time mu sync.RWMutex + // dialMu serializes dials to this one node. It is never held together + // with Router.mu, so a slow dial to one node cannot stall routing to any + // other node (see getClient). + dialMu sync.Mutex +} + +// closeConn closes the node's current connection, if any. +func (s *clientState) closeConn() error { + s.mu.RLock() + conn := s.conn + s.mu.RUnlock() + if conn == nil { + return nil + } + return conn.Close() } // ConnectionStats represents statistics for a single connection @@ -127,13 +142,25 @@ func NewRouterWithConfig(ring Ring, localID string, config *RouterConfig) *Route // Route returns a client for routing requests for the given key // Returns an error if the key should be handled locally (defensive check) func (r *Router) Route(key string) (pb.CacheServiceClient, error) { - return r.RouteWithRetry(key, r.config.MaxRetries) + return r.RouteContext(context.Background(), key) +} + +// RouteContext is Route bound to the caller's context. A cancelled or expired +// ctx ends the retry backoff and any in-flight dial promptly, so an abandoned +// request does not keep dialing and backing off against the owner on the +// caller's behalf. +func (r *Router) RouteContext(ctx context.Context, key string) (pb.CacheServiceClient, error) { + return r.routeWithRetry(ctx, key, r.config.MaxRetries) } // RouteWithRetry returns a client for routing with configurable retry attempts // Returns an error if the key maps to the local node (this should not happen // as callers should check IsLocal first, but we check defensively) func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClient, error) { + return r.routeWithRetry(context.Background(), key, maxRetries) +} + +func (r *Router) routeWithRetry(ctx context.Context, key string, maxRetries int) (pb.CacheServiceClient, error) { node, err := r.ring.GetNode(key) if err != nil { metrics.ClusterRouteRequests.WithLabelValues("error").Inc() @@ -160,8 +187,15 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie if attempt > 0 { // Wait before retry with jittered exponential backoff. Jitter spreads // retries so a degraded ring's survivors don't all reconnect to a - // recovered node in lockstep (thundering herd, issue #164). - time.Sleep(jittered(backoff)) + // recovered node in lockstep (thundering herd, issue #164). The wait + // ends early if the caller gives up. + timer := time.NewTimer(jittered(backoff)) + select { + case <-ctx.Done(): + timer.Stop() + return nil, ctx.Err() + case <-timer.C: + } backoff = r.calculateBackoff(backoff) metrics.ClusterRetryAttempts.WithLabelValues(node.ID).Inc() @@ -173,7 +207,7 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie Msg("Retrying connection after failure") } - client, err := r.getClient(node.ID) + client, err := r.getClient(ctx, node.ID) if err == nil { zlog.Debug(). Str("node_id", node.ID). @@ -185,6 +219,12 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie lastErr = err + // The caller gave up: report its own context error rather than + // spending the retry budget it no longer wants. + if ctx.Err() != nil { + return nil, ctx.Err() + } + // Don't retry if it's not a retryable error if !IsRetryableError(err) { break @@ -196,8 +236,16 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie return nil, NewMaxRetriesExceededError(node.ID, key, maxRetries+1, lastErr) } -// getClient returns a client for the given node, creating one if necessary -func (r *Router) getClient(nodeID string) (pb.CacheServiceClient, error) { +// getClient returns a client for the given node, creating one if necessary. +// +// Router.mu guards only the clients map and is never held across a dial. A +// blocking dial can take up to ConnectionTimeout, and holding the router-wide +// write lock for that long stalled every route to every other node behind one +// unreachable peer, so a single restarting node froze all forwarding on the +// cluster. Dials to one node are serialized on that node's own dialMu +// instead: concurrent callers for the same node share one attempt, and callers +// for other nodes are unaffected. +func (r *Router) getClient(ctx context.Context, nodeID string) (pb.CacheServiceClient, error) { // Fast path: check if client exists and is healthy r.mu.RLock() state, exists := r.clients[nodeID] @@ -208,30 +256,39 @@ func (r *Router) getClient(nodeID string) (pb.CacheServiceClient, error) { Msg("Checking if client exists and is healthy") if exists && state != nil { - if err := r.getConnectionHealth(state, nodeID); err == nil { - return state.client, nil + client, err := r.healthyClient(state, nodeID) + if err == nil { + return client, nil } else if errors.Is(err, ErrCircuitBreakerOpen) { metrics.ClusterRoutingErrors.WithLabelValues("circuit_breaker_open").Inc() return nil, err } } - // Slow path: create new client or reconnect + // Slow path: create new client or reconnect. Register the node's state + // under Router.mu, then dial under the node's own lock only. r.mu.Lock() - defer r.mu.Unlock() + state = r.clients[nodeID] + if state == nil { + state = &clientState{} + r.clients[nodeID] = state + } + r.mu.Unlock() + + state.dialMu.Lock() + defer state.dialMu.Unlock() zlog.Debug(). Str("node_id", nodeID). Msg("Creating new client or reconnecting") - // Double-check after acquiring write lock - state, exists = r.clients[nodeID] - if exists && state != nil { - if err := r.getConnectionHealth(state, nodeID); err == nil { - return state.client, nil - } else if errors.Is(err, ErrCircuitBreakerOpen) { - return nil, err - } + // Double-check after acquiring the dial lock: a concurrent caller for the + // same node may have just connected. + client, err := r.healthyClient(state, nodeID) + if err == nil { + return client, nil + } else if errors.Is(err, ErrCircuitBreakerOpen) { + return nil, err } // Get node listen address from ring @@ -254,20 +311,14 @@ func (r *Router) getClient(nodeID string) (pb.CacheServiceClient, error) { return nil, NewNodeNotFoundError(nodeID, "") } - // Create or update client state - if state == nil { - state = &clientState{} - r.clients[nodeID] = state - } - - // Close existing connection if any - if state.conn != nil { - state.conn.Close() - } - // Create new connection with keepalive - conn, err := r.createConnection(nodeAddr) + conn, err := r.createConnection(ctx, nodeAddr) if err != nil { + if ctx.Err() != nil { + // The caller gave up mid-dial; that says nothing about the node, + // so it must not count toward its circuit breaker. + return nil, ctx.Err() + } r.recordFailureAndOpenCircuit(state, nodeID) logsample.DegradedRing(). @@ -280,9 +331,27 @@ func (r *Router) getClient(nodeID string) (pb.CacheServiceClient, error) { return nil, NewConnectionFailedError(nodeID, nodeAddr, err) } - client := pb.NewCacheServiceClient(conn) + client = pb.NewCacheServiceClient(conn) + state.mu.Lock() + old := state.conn state.client = client state.conn = conn + state.mu.Unlock() + if old != nil { + old.Close() + } + + // RemoveClient or RefreshConnections may have dropped this node while the + // dial was in flight. They closed the connection they could see, so close + // this one too rather than hand out a connection nothing will ever close. + r.mu.RLock() + current := r.clients[nodeID] + r.mu.RUnlock() + if current != state { + conn.Close() + metrics.ClusterRoutingErrors.WithLabelValues("node_not_found").Inc() + return nil, NewNodeNotFoundError(nodeID, nodeAddr) + } atomic.StoreInt32(&state.failureCount, 0) // Reset failure count on successful connection metrics.ClusterConnectionsActive.WithLabelValues(nodeID).Set(1) @@ -295,9 +364,21 @@ func (r *Router) getClient(nodeID string) (pb.CacheServiceClient, error) { return client, nil } +// healthyClient returns the node's client when its connection is usable, or +// the reason it is not. +func (r *Router) healthyClient(state *clientState, nodeID string) (pb.CacheServiceClient, error) { + if err := r.getConnectionHealth(state, nodeID); err != nil { + return nil, err + } + state.mu.RLock() + client := state.client + state.mu.RUnlock() + return client, nil +} + // createConnection creates a new gRPC connection with configured parameters -func (r *Router) createConnection(address string) (*grpc.ClientConn, error) { - ctx, cancel := context.WithTimeout(context.Background(), r.config.ConnectionTimeout) +func (r *Router) createConnection(ctx context.Context, address string) (*grpc.ClientConn, error) { + ctx, cancel := context.WithTimeout(ctx, r.config.ConnectionTimeout) defer cancel() zlog.Debug(). @@ -346,9 +427,12 @@ func (r *Router) getConnectionHealth(state *clientState, nodeID string) error { } // Check connection state + state.mu.RLock() + conn := state.conn + state.mu.RUnlock() error := fmt.Errorf("client state not set for node %s", nodeID) - if state.conn != nil { - connState := state.conn.GetState() + if conn != nil { + connState := conn.GetState() if connState != connectivity.Shutdown && connState != connectivity.TransientFailure { return nil } @@ -444,7 +528,7 @@ func (r *Router) GetClientForNode(nodeID string) (pb.CacheServiceClient, error) return nil, NewLocalRoutingError(r.localID, "") } - return r.getClient(nodeID) + return r.getClient(context.Background(), nodeID) } // RemoveClient removes and closes the client connection for a node @@ -453,9 +537,7 @@ func (r *Router) RemoveClient(nodeID string) { defer r.mu.Unlock() if state, exists := r.clients[nodeID]; exists { - if state.conn != nil { - state.conn.Close() - } + state.closeConn() delete(r.clients, nodeID) metrics.ClusterConnectionsActive.WithLabelValues(nodeID).Set(0) @@ -471,13 +553,14 @@ func (r *Router) Close() error { defer r.mu.Unlock() for nodeID, state := range r.clients { - if state != nil && state.conn != nil { - if err := state.conn.Close(); err != nil { - zlog.Error(). - Err(err). - Str("node_id", nodeID). - Msg("Error closing connection") - } + if state == nil { + continue + } + if err := state.closeConn(); err != nil { + zlog.Error(). + Err(err). + Str("node_id", nodeID). + Msg("Error closing connection") } } @@ -507,8 +590,8 @@ func (r *Router) RefreshConnections() { // Remove connections to inactive nodes for nodeID, state := range r.clients { if !activeNodeMap[nodeID] { - if state != nil && state.conn != nil { - state.conn.Close() + if state != nil { + state.closeConn() } delete(r.clients, nodeID) @@ -527,12 +610,11 @@ func (r *Router) GetConnectionStats() map[string]ConnectionStats { continue } + state.mu.RLock() var connState connectivity.State if state.conn != nil { connState = state.conn.GetState() } - - state.mu.RLock() lastFailure := state.lastFailure circuitOpenTime := state.circuitOpenTime state.mu.RUnlock() diff --git a/coordinator/router_dial_test.go b/coordinator/router_dial_test.go new file mode 100644 index 00000000..3f9ff86b --- /dev/null +++ b/coordinator/router_dial_test.go @@ -0,0 +1,134 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +package coordinator + +import ( + "context" + "net" + "testing" + "time" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "google.golang.org/grpc" +) + +// blockingDialer never completes a connection to blockAddr: it holds the dial +// until its context ends, the shape of a peer that is mid-restart and +// blackholing SYNs. Every other address connects normally. +func blockingDialer(blockAddr string) grpc.DialOption { + return grpc.WithContextDialer(func(ctx context.Context, addr string) (net.Conn, error) { + if addr == blockAddr { + <-ctx.Done() + return nil, ctx.Err() + } + var d net.Dialer + return d.DialContext(ctx, "tcp", addr) + }) +} + +func dialTestConfig(connTimeout time.Duration, opts ...grpc.DialOption) *RouterConfig { + cfg := DefaultRouterConfig() + cfg.ConnectionTimeout = connTimeout + cfg.MaxRetries = 0 + cfg.GRPCDialOptions = opts + return cfg +} + +const stuckAddr = "stuck.invalid:1" + +// TestRouter_SlowDialDoesNotBlockOtherNodes pins the head-of-line fix: while a +// dial to one unreachable node is in flight, routing to a healthy node must +// complete immediately instead of waiting out that dial's timeout behind the +// router-wide lock. +func TestRouter_SlowDialDoesNotBlockOtherNodes(t *testing.T) { + mockRing := newMockRing("local-node") + cacheAddr, clusterAddr, clusterSrv, cacheSrv, _, _ := startMockRouterServer(t, "healthy-node") + defer clusterSrv.Stop() + defer cacheSrv.Stop() + mockRing.AddNode("healthy-node", clusterAddr, cacheAddr) + mockRing.AddNode("stuck-node", "localhost:1", stuckAddr) + mockRing.SetKeyOwner("healthy-key", "healthy-node") + mockRing.SetKeyOwner("stuck-key", "stuck-node") + + const connTimeout = 3 * time.Second + router := NewRouterWithConfig(mockRing, "local-node", dialTestConfig(connTimeout, blockingDialer(stuckAddr))) + defer router.Close() + + stuckDone := make(chan error, 1) + go func() { + _, err := router.Route("stuck-key") + stuckDone <- err + }() + // The stuck node's state is registered before its dial starts, so its + // appearance in the stats means the dial is now in flight. + require.Eventually(t, func() bool { + _, ok := router.GetConnectionStats()["stuck-node"] + return ok + }, time.Second, 5*time.Millisecond, "stuck dial never started") + + start := time.Now() + client, err := router.Route("healthy-key") + require.NoError(t, err) + require.NotNil(t, client) + assert.Less(t, time.Since(start), connTimeout/2, + "routing to a healthy node waited on another node's dial") + + select { + case err := <-stuckDone: + assert.Error(t, err) + case <-time.After(2 * connTimeout): + t.Fatal("stuck dial never timed out") + } +} + +// TestRouter_CancelledContextAbortsDial: a caller that gives up mid-dial gets +// its own context error back promptly, and the abandoned dial is not held +// against the node's circuit breaker. +func TestRouter_CancelledContextAbortsDial(t *testing.T) { + mockRing := newMockRing("local-node") + mockRing.AddNode("stuck-node", "localhost:1", stuckAddr) + mockRing.SetKeyOwner("stuck-key", "stuck-node") + + cfg := dialTestConfig(5*time.Second, blockingDialer(stuckAddr)) + cfg.CircuitBreakerThreshold = 1 + router := NewRouterWithConfig(mockRing, "local-node", cfg) + defer router.Close() + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + start := time.Now() + _, err := router.RouteContext(ctx, "stuck-key") + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(start), 2*time.Second, "dial outlived the caller's context") + + stats := router.GetConnectionStats()["stuck-node"] + assert.Equal(t, int32(0), stats.FailureCount, "an abandoned dial is not a node failure") + assert.False(t, stats.CircuitOpen) +} + +// TestRouter_CancelledContextStopsRetryBackoff: the retry sleep between +// attempts ends when the caller's context does. +func TestRouter_CancelledContextStopsRetryBackoff(t *testing.T) { + mockRing := newMockRing("local-node") + mockRing.AddNode("dead-node", "localhost:1", "localhost:1") // nothing listens here + mockRing.SetKeyOwner("dead-key", "dead-node") + + cfg := dialTestConfig(100 * time.Millisecond) + cfg.MaxRetries = 3 + cfg.InitialRetryBackoff = 10 * time.Second + cfg.MaxRetryBackoff = 10 * time.Second + router := NewRouterWithConfig(mockRing, "local-node", cfg) + defer router.Close() + + ctx, cancel := context.WithCancel(context.Background()) + go func() { + time.Sleep(300 * time.Millisecond) + cancel() + }() + start := time.Now() + _, err := router.RouteContext(ctx, "dead-key") + assert.ErrorIs(t, err, context.Canceled) + assert.Less(t, time.Since(start), 2*time.Second, "retry backoff outlived the caller's context") +} diff --git a/server/operations/cas.go b/server/operations/cas.go index 2e119070..ec15931a 100644 --- a/server/operations/cas.go +++ b/server/operations/cas.go @@ -103,7 +103,7 @@ func (o *Operations) putStreamIfVersionRemote(ctx context.Context, key string, r if err != nil { return 0, err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { return 0, err } @@ -157,7 +157,7 @@ func (o *Operations) getStreamWithVersionRemote(ctx context.Context, key string, if err != nil { return 0, false, err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { return 0, false, err } @@ -228,7 +228,7 @@ func (o *Operations) getWithVersionRemote(ctx context.Context, key string) ([]by logsample.DegradedRing().Err(err).Str("key", key).Msg("Hop count limit exceeded for GetWithVersion") return nil, 0, false, err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { logsample.DegradedRing().Err(err).Str("key", key).Msg("Failed to route key for GetWithVersion") return nil, 0, false, err @@ -249,7 +249,7 @@ func (o *Operations) putIfVersionRemote(ctx context.Context, key string, data [] logsample.DegradedRing().Err(err).Str("key", key).Msg("Hop count limit exceeded for PutIfVersion") return 0, err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { logsample.DegradedRing().Err(err).Str("key", key).Msg("Failed to route key for PutIfVersion") return 0, err @@ -278,7 +278,7 @@ func (o *Operations) deleteIfVersionRemote(ctx context.Context, key string, expe logsample.DegradedRing().Err(err).Str("key", key).Msg("Hop count limit exceeded for DeleteIfVersion") return err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { logsample.DegradedRing().Err(err).Str("key", key).Msg("Failed to route key for DeleteIfVersion") return err diff --git a/server/operations/delete.go b/server/operations/delete.go index b046c257..f2635cff 100644 --- a/server/operations/delete.go +++ b/server/operations/delete.go @@ -49,7 +49,7 @@ func (o *Operations) deleteRemote(ctx context.Context, key string) error { return err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { logsample.DegradedRing().Err(err).Str("key", key).Msg("Failed to route key for delete") return err diff --git a/server/operations/get.go b/server/operations/get.go index afb9171a..6cd32315 100644 --- a/server/operations/get.go +++ b/server/operations/get.go @@ -84,7 +84,7 @@ func (o *Operations) getRemote(ctx context.Context, key string, start, end int64 return nil, false, err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { logsample.DegradedRing().Err(err).Str("key", key).Msg("Failed to route key") return nil, false, err diff --git a/server/operations/operations.go b/server/operations/operations.go index e777f788..f0c8b124 100644 --- a/server/operations/operations.go +++ b/server/operations/operations.go @@ -6,6 +6,8 @@ package operations import ( + "context" + "github.com/tigrisdata/ocache/coordinator" "github.com/tigrisdata/ocache/coordinator/ring" pb "github.com/tigrisdata/ocache/proto" @@ -46,6 +48,15 @@ func (o *Operations) Route(key string) (pb.CacheServiceClient, error) { return o.coordinator.Route(key) } +// RouteContext is Route bound to the caller's context, so an abandoned request +// stops dialing and backing off against the owner node promptly. +func (o *Operations) RouteContext(ctx context.Context, key string) (pb.CacheServiceClient, error) { + if o.coordinator == nil { + return nil, nil + } + return o.coordinator.RouteContext(ctx, key) +} + // IsClusterMode returns true if clustering is enabled. func (o *Operations) IsClusterMode() bool { return o.coordinator != nil diff --git a/server/operations/put.go b/server/operations/put.go index 6e63c1c6..5b188140 100644 --- a/server/operations/put.go +++ b/server/operations/put.go @@ -72,7 +72,7 @@ func (o *Operations) putRemote(ctx context.Context, key string, body io.Reader, return err } - client, err := o.Route(key) + client, err := o.RouteContext(ctx, key) if err != nil { logsample.DegradedRing().Err(err).Str("key", key).Msg("Failed to route key for put") return err diff --git a/server/service/cas.go b/server/service/cas.go index 5de25df3..d4a92069 100644 --- a/server/service/cas.go +++ b/server/service/cas.go @@ -290,7 +290,7 @@ func (s *CacheService) GetStreamWithVersion(req *pb.GetRequest, stream pb.CacheS // forwardStreamingPutIfVersion forwards a streaming CAS put to the owner node. func (s *CacheService) forwardStreamingPutIfVersion(localStream pb.CacheService_PutStreamIfVersionServer, first *pb.PutIfVersionRequest) error { - client, err := s.coordinator.Route(first.Key) + client, err := s.coordinator.RouteContext(localStream.Context(), first.Key) if err != nil { return status.Error(codes.Unavailable, fmt.Sprintf("routing error: %v", err)) } @@ -336,7 +336,7 @@ func (s *CacheService) forwardStreamingPutIfVersion(localStream pb.CacheService_ // forwardStreamingGetWithVersion forwards a streaming CAS read from the owner. func (s *CacheService) forwardStreamingGetWithVersion(req *pb.GetRequest, localStream pb.CacheService_GetStreamWithVersionServer) error { - client, err := s.coordinator.Route(req.Key) + client, err := s.coordinator.RouteContext(localStream.Context(), req.Key) if err != nil { return status.Error(codes.Unavailable, fmt.Sprintf("routing error: %v", err)) } diff --git a/server/service/get.go b/server/service/get.go index 39238cbc..33bf106a 100644 --- a/server/service/get.go +++ b/server/service/get.go @@ -80,7 +80,7 @@ func (s *CacheService) Get(req *pb.GetRequest, stream pb.CacheService_GetServer) // forwardStreamingGet forwards a streaming Get request to a remote node func (s *CacheService) forwardStreamingGet(req *pb.GetRequest, localStream pb.CacheService_GetServer) error { // Forward to the correct node - client, err := s.coordinator.Route(req.Key) + client, err := s.coordinator.RouteContext(localStream.Context(), req.Key) if err != nil { metrics.RPCRequests.WithLabelValues("Get", "routing_error").Inc() return status.Errorf(codes.Unavailable, "routing error: %v", err) diff --git a/server/service/put.go b/server/service/put.go index 0d6b48e8..6364f6ee 100644 --- a/server/service/put.go +++ b/server/service/put.go @@ -124,7 +124,7 @@ func (s *CacheService) forwardStreamingPut(localStream pb.CacheService_PutServer key := firstChunk.Key // Forward to the correct node - client, err := s.coordinator.Route(key) + client, err := s.coordinator.RouteContext(localStream.Context(), key) if err != nil { metrics.RPCRequests.WithLabelValues("Put", "routing_error").Inc() return localStream.SendAndClose(&pb.PutResponse{ From b036fb7d6323b555e09ec3ec4ea7544ad70f5b8e Mon Sep 17 00:00:00 2001 From: Ovais Tariq Date: Fri, 25 Sep 2026 20:34:23 -0700 Subject: [PATCH 2/2] fix(router): let waiters give up on a stalled dial; dial List peers under ctx Review follow-ups on the dial-outside-the-lock change. A second caller for a node whose dial was already in flight blocked on the node's dial mutex without regard to its own deadline, so it could wait out the full connection timeout behind the first caller. The per-node lock is now a one-slot semaphore taken with a select on ctx.Done, so a waiter returns its own context error as soon as its deadline passes. GetClientForNode, which the cluster-wide List fan-out uses, now takes the request context so a cancelled List does not wait for a blackholed peer's dial either. Co-Authored-By: Claude Fable 5.1 Claude-Session: https://claude.ai/code/session_01YCoN4Him5ZY4hZovVp37ra Signed-off-by: Ovais Tariq --- coordinator/router.go | 31 +++++++++++++++++++-------- coordinator/router_dial_test.go | 37 +++++++++++++++++++++++++++++++++ server/operations/list.go | 2 +- 3 files changed, 60 insertions(+), 10 deletions(-) diff --git a/coordinator/router.go b/coordinator/router.go index 1c56493a..c7925bfc 100644 --- a/coordinator/router.go +++ b/coordinator/router.go @@ -86,10 +86,15 @@ type clientState struct { circuitOpen int32 // atomic: 0=closed, 1=open lastFailure time.Time mu sync.RWMutex - // dialMu serializes dials to this one node. It is never held together + // dialSem serializes dials to this one node (a one-slot semaphore, so a + // waiter can give up when its context ends). It is never held together // with Router.mu, so a slow dial to one node cannot stall routing to any // other node (see getClient). - dialMu sync.Mutex + dialSem chan struct{} +} + +func newClientState() *clientState { + return &clientState{dialSem: make(chan struct{}, 1)} } // closeConn closes the node's current connection, if any. @@ -270,13 +275,20 @@ func (r *Router) getClient(ctx context.Context, nodeID string) (pb.CacheServiceC r.mu.Lock() state = r.clients[nodeID] if state == nil { - state = &clientState{} + state = newClientState() r.clients[nodeID] = state } r.mu.Unlock() - state.dialMu.Lock() - defer state.dialMu.Unlock() + // Take the node's dial slot, or give up with the caller: a second request + // for a node whose dial is stalled must not outlive its own deadline + // waiting for that dial. + select { + case state.dialSem <- struct{}{}: + defer func() { <-state.dialSem }() + case <-ctx.Done(): + return nil, ctx.Err() + } zlog.Debug(). Str("node_id", nodeID). @@ -520,15 +532,16 @@ func (r *Router) IsLocal(key string) bool { return r.ring.IsLocal(key) } -// GetClientForNode returns a client for a specific node ID -// This is useful for operations that need to query all nodes (e.g., List) -func (r *Router) GetClientForNode(nodeID string) (pb.CacheServiceClient, error) { +// GetClientForNode returns a client for a specific node ID, dialing it under +// ctx if needed. This is useful for operations that need to query all nodes +// (e.g., List). +func (r *Router) GetClientForNode(ctx context.Context, nodeID string) (pb.CacheServiceClient, error) { // Check if this is the local node (shouldn't be called for local, but check defensively) if nodeID == r.localID { return nil, NewLocalRoutingError(r.localID, "") } - return r.getClient(context.Background(), nodeID) + return r.getClient(ctx, nodeID) } // RemoveClient removes and closes the client connection for a node diff --git a/coordinator/router_dial_test.go b/coordinator/router_dial_test.go index 3f9ff86b..0acec667 100644 --- a/coordinator/router_dial_test.go +++ b/coordinator/router_dial_test.go @@ -108,6 +108,43 @@ func TestRouter_CancelledContextAbortsDial(t *testing.T) { assert.False(t, stats.CircuitOpen) } +// TestRouter_CancelledContextStopsWaitingForDial: a second caller for a node +// whose dial is already stalled gives up when its own context ends instead +// of queueing behind that dial for the full connection timeout. +func TestRouter_CancelledContextStopsWaitingForDial(t *testing.T) { + mockRing := newMockRing("local-node") + mockRing.AddNode("stuck-node", "localhost:1", stuckAddr) + mockRing.SetKeyOwner("stuck-key", "stuck-node") + + const connTimeout = 3 * time.Second + router := NewRouterWithConfig(mockRing, "local-node", dialTestConfig(connTimeout, blockingDialer(stuckAddr))) + defer router.Close() + + firstDone := make(chan error, 1) + go func() { + _, err := router.Route("stuck-key") + firstDone <- err + }() + require.Eventually(t, func() bool { + _, ok := router.GetConnectionStats()["stuck-node"] + return ok + }, time.Second, 5*time.Millisecond, "stuck dial never started") + + ctx, cancel := context.WithTimeout(context.Background(), 200*time.Millisecond) + defer cancel() + start := time.Now() + _, err := router.RouteContext(ctx, "stuck-key") + assert.ErrorIs(t, err, context.DeadlineExceeded) + assert.Less(t, time.Since(start), connTimeout/2, "second caller waited on the first caller's dial") + + select { + case err := <-firstDone: + assert.Error(t, err) + case <-time.After(2 * connTimeout): + t.Fatal("stuck dial never timed out") + } +} + // TestRouter_CancelledContextStopsRetryBackoff: the retry sleep between // attempts ends when the caller's context does. func TestRouter_CancelledContextStopsRetryBackoff(t *testing.T) { diff --git a/server/operations/list.go b/server/operations/list.go index bea5af6b..d4067ba6 100644 --- a/server/operations/list.go +++ b/server/operations/list.go @@ -321,7 +321,7 @@ func (o *Operations) fetchFromAllNodes(ctx context.Context, nodes []*ring.NodeIn resp.Keys, resp.ContinuationToken, resp.HasMore, err = o.ListLocal(ctx, prefix, limit, startKey) } } else { - client, clientErr := router.GetClientForNode(n.ID) + client, clientErr := router.GetClientForNode(ctx, n.ID) if clientErr != nil { logsample.DegradedRing().Err(clientErr).Str("node_id", n.ID).Msg("Failed to get client for node, skipping") return