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..c7925bfc 100644 --- a/coordinator/router.go +++ b/coordinator/router.go @@ -79,13 +79,33 @@ 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 + // 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). + dialSem chan struct{} +} + +func newClientState() *clientState { + return &clientState{dialSem: make(chan struct{}, 1)} +} + +// 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 +147,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 +192,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 +212,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 +224,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 +241,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 +261,46 @@ 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 = newClientState() + r.clients[nodeID] = state + } + r.mu.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). 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 +323,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 +343,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 +376,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 +439,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 } @@ -436,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(nodeID) + return r.getClient(ctx, nodeID) } // RemoveClient removes and closes the client connection for a node @@ -453,9 +550,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 +566,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 +603,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 +623,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..0acec667 --- /dev/null +++ b/coordinator/router_dial_test.go @@ -0,0 +1,171 @@ +// 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_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) { + 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/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 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{