diff --git a/coordinator/coordinator.go b/coordinator/coordinator.go index b3b066db..b264c999 100644 --- a/coordinator/coordinator.go +++ b/coordinator/coordinator.go @@ -340,6 +340,17 @@ func (c *Coordinator) Route(key string) (pb.CacheServiceClient, error) { return c.router.Route(key) } +// Resolve returns the current owner for a key without acquiring a peer client. +func (c *Coordinator) Resolve(key string) (*ring.NodeInfo, error) { + return c.router.Resolve(key) +} + +// RouteToNode returns a client for an already-resolved owner while preserving +// the router's connection retry, health, and circuit-breaker behavior. +func (c *Coordinator) RouteToNode(nodeID string) (pb.CacheServiceClient, error) { + return c.router.RouteToNode(nodeID) +} + // 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/interceptors.go b/coordinator/interceptors.go index 143e26db..fb6491df 100644 --- a/coordinator/interceptors.go +++ b/coordinator/interceptors.go @@ -59,9 +59,12 @@ type ResponseMetadata struct { ForwardedBy string } -// ExtractRequestMetadata extracts routing metadata from the incoming context +// ExtractRequestMetadata extracts routing metadata from incoming or outgoing context. func ExtractRequestMetadata(ctx context.Context) RequestMetadata { md, ok := metadata.FromIncomingContext(ctx) + if !ok { + md, ok = metadata.FromOutgoingContext(ctx) + } if !ok { return RequestMetadata{} } @@ -97,12 +100,17 @@ func ExtractRequestMetadata(ctx context.Context) RequestMetadata { // AttachForwardingMetadata attaches forwarding metadata to outgoing context func AttachForwardingMetadata(ctx context.Context, rm RequestMetadata) context.Context { - return metadata.AppendToOutgoingContext(ctx, - MetadataKeyRingEpoch, strconv.FormatUint(rm.RingEpoch, 10), - MetadataKeyHop, strconv.Itoa(rm.HopCount), - MetadataKeyForwarded, strconv.FormatBool(rm.Forwarded), - MetadataKeyOrigin, rm.OriginNode, - ) + md, ok := metadata.FromOutgoingContext(ctx) + if !ok { + md = metadata.MD{} + } else { + md = md.Copy() + } + md.Set(MetadataKeyRingEpoch, strconv.FormatUint(rm.RingEpoch, 10)) + md.Set(MetadataKeyHop, strconv.Itoa(rm.HopCount)) + md.Set(MetadataKeyForwarded, strconv.FormatBool(rm.Forwarded)) + md.Set(MetadataKeyOrigin, rm.OriginNode) + return metadata.NewOutgoingContext(ctx, md) } // CheckHopCount validates that the hop count hasn't exceeded MaxHops. diff --git a/coordinator/interceptors_test.go b/coordinator/interceptors_test.go index ea96ab9a..4fc86341 100644 --- a/coordinator/interceptors_test.go +++ b/coordinator/interceptors_test.go @@ -105,6 +105,26 @@ func TestAttachForwardingMetadata(t *testing.T) { assert.Equal(t, []string{"node1"}, md.Get(MetadataKeyOrigin)) } +func TestIncrementHopCountUsesOutgoingMetadata(t *testing.T) { + ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs( + MetadataKeyRingEpoch, "42", + MetadataKeyHop, "1", + MetadataKeyOrigin, "origin", + "authorization", "bearer", + )) + + forwarded, err := IncrementHopCount(ctx, "gateway") + require.NoError(t, err) + + md, ok := metadata.FromOutgoingContext(forwarded) + require.True(t, ok) + assert.Equal(t, []string{"42"}, md.Get(MetadataKeyRingEpoch)) + assert.Equal(t, []string{"2"}, md.Get(MetadataKeyHop)) + assert.Equal(t, []string{"true"}, md.Get(MetadataKeyForwarded)) + assert.Equal(t, []string{"origin"}, md.Get(MetadataKeyOrigin)) + assert.Equal(t, []string{"bearer"}, md.Get("authorization")) +} + func TestCheckHopCount(t *testing.T) { tests := []struct { name string diff --git a/coordinator/router.go b/coordinator/router.go index 83fccaed..c4a637f5 100644 --- a/coordinator/router.go +++ b/coordinator/router.go @@ -124,6 +124,13 @@ func NewRouterWithConfig(ring Ring, localID string, config *RouterConfig) *Route } } +// Resolve returns the current owner for a key without acquiring a peer client. +// Keeping ownership lookup separate lets callers that already resolved a key +// acquire the corresponding cached client without hashing the key again. +func (r *Router) Resolve(key string) (*ring.NodeInfo, error) { + return r.ring.GetNode(key) +} + // 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) { @@ -134,7 +141,7 @@ func (r *Router) Route(key string) (pb.CacheServiceClient, error) { // 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) { - node, err := r.ring.GetNode(key) + node, err := r.Resolve(key) if err != nil { metrics.ClusterRouteRequests.WithLabelValues("error").Inc() return nil, err @@ -153,6 +160,17 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie return nil, NewLocalRoutingError(r.localID, key) } + return r.routeToNode(node.ID, key, maxRetries) +} + +// RouteToNode returns a client for an already-resolved owner. It uses the same +// connection retry, health, and circuit-breaker path as Route, but does not +// resolve the key again. +func (r *Router) RouteToNode(nodeID string) (pb.CacheServiceClient, error) { + return r.routeToNode(nodeID, "", r.config.MaxRetries) +} + +func (r *Router) routeToNode(nodeID, key string, maxRetries int) (pb.CacheServiceClient, error) { var lastErr error backoff := r.config.InitialRetryBackoff @@ -164,21 +182,17 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie time.Sleep(jittered(backoff)) backoff = r.calculateBackoff(backoff) - metrics.ClusterRetryAttempts.WithLabelValues(node.ID).Inc() + metrics.ClusterRetryAttempts.WithLabelValues(nodeID).Inc() zlog.Debug(). - Str("node_id", node.ID). + Str("node_id", nodeID). Int("attempt", attempt). Dur("backoff", backoff). Msg("Retrying connection after failure") } - client, err := r.getClient(node.ID) + client, err := r.getClient(nodeID) if err == nil { - zlog.Debug(). - Str("node_id", node.ID). - Msg("Successfully routed to node") - metrics.ClusterRouteRequests.WithLabelValues("remote").Inc() return client, nil } @@ -193,7 +207,7 @@ func (r *Router) RouteWithRetry(key string, maxRetries int) (pb.CacheServiceClie metrics.ClusterRouteRequests.WithLabelValues("error").Inc() metrics.ClusterRoutingErrors.WithLabelValues("max_retries_exceeded").Inc() - return nil, NewMaxRetriesExceededError(node.ID, key, maxRetries+1, lastErr) + return nil, NewMaxRetriesExceededError(nodeID, key, maxRetries+1, lastErr) } // getClient returns a client for the given node, creating one if necessary @@ -291,6 +305,9 @@ func (r *Router) getClient(nodeID string) (pb.CacheServiceClient, error) { Str("node_id", nodeID). Str("address", nodeAddr). Msg("Created connection to node") + zlog.Debug(). + Str("node_id", nodeID). + Msg("Successfully routed to node") return client, nil } diff --git a/coordinator/router_test.go b/coordinator/router_test.go index 09cda18a..61e19ac8 100644 --- a/coordinator/router_test.go +++ b/coordinator/router_test.go @@ -309,6 +309,15 @@ func TestRouter_BasicRouting(t *testing.T) { client, err = router.Route("remote-key") assert.NoError(t, err) assert.NotNil(t, client, "Remote routing should return a client") + + // An already-resolved owner should reuse the same client acquisition path + // without resolving the key again. + owner, err := router.Resolve("remote-key") + require.NoError(t, err) + assert.Equal(t, "remote-node", owner.ID) + resolvedClient, err := router.RouteToNode(owner.ID) + assert.NoError(t, err) + assert.NotNil(t, resolvedClient) } func TestRouter_CircuitBreaker(t *testing.T) { diff --git a/scripts/benchmark_http_gateway.py b/scripts/benchmark_http_gateway.py new file mode 100755 index 00000000..67d6311e --- /dev/null +++ b/scripts/benchmark_http_gateway.py @@ -0,0 +1,296 @@ +#!/usr/bin/env python3 +"""Measure HTTP PutObject forwarding through a warmed three-node cluster.""" + +import argparse +import base64 +import http.client +import json +import pathlib +import socket +import subprocess +import sys +import tempfile +import time +import urllib.parse + + +HOST = "127.0.0.1" + + +def free_ports(count): + sockets = [] + ports = [] + try: + for _ in range(count): + sock = socket.socket() + sock.bind((HOST, 0)) + sockets.append(sock) + ports.append(sock.getsockname()[1]) + finally: + for sock in sockets: + sock.close() + return ports + + +def fetch_json(port, path): + conn = http.client.HTTPConnection(HOST, port, timeout=2) + try: + conn.request("GET", path) + response = conn.getresponse() + body = response.read() + if response.status != 200: + raise RuntimeError(f"GET {path} on {port}: HTTP {response.status}: {body[:200]!r}") + return json.loads(body) + finally: + conn.close() + + +def wait_for_http(port, path, predicate=None, timeout=60): + deadline = time.monotonic() + timeout + last_error = None + while time.monotonic() < deadline: + try: + value = fetch_json(port, path) + if predicate is None or predicate(value): + return value + except (OSError, RuntimeError, ValueError) as error: + last_error = error + time.sleep(0.1) + raise RuntimeError(f"timed out waiting for http://{HOST}:{port}{path}: {last_error}") + + +def active_nodes(topology): + payload = topology.get("topology", topology) + nodes = payload.get("nodes", []) + return [node for node in nodes if "ACTIVE" in str(node.get("status", "")) or node.get("status") == 0] + + +def topology_ready(value): + return len(active_nodes(value)) == 3 + + +def node_id(node): + return node.get("id", node.get("nodeId", "")) + + +def node_tokens(topology): + payload = topology.get("topology", topology) + ring_config = payload.get("ringConfig", payload.get("ring_config", {})) + assignments = ring_config.get("nodeTokens", ring_config.get("node_tokens", [])) + tokens = [] + for assignment in assignments: + owner = assignment.get("nodeId", assignment.get("node_id", "")) + for token in assignment.get("tokens", []): + tokens.append((int(token), owner)) + if not tokens: + raise RuntimeError("topology did not contain ring tokens") + return sorted(tokens) + + +def fnv1a32(value): + result = 2166136261 + for byte in value.encode(): + result ^= byte + result = (result * 16777619) & 0xFFFFFFFF + return result + + +def owner_for(key, tokens): + token = fnv1a32(key) + for ring_token, owner in tokens: + if ring_token >= token: + return owner + return tokens[0][1] + + +def scalar_metric(text, name): + for line in text.splitlines(): + if line.startswith(name + " "): + return float(line.split()[1]) + return 0.0 + + +def metric_value(text, method, status): + metric = "ocache_rpc_requests_total" + wanted = {"method": method, "status": status} + for line in text.splitlines(): + if not line.startswith(metric + "{"): + continue + labels, _, value = line.partition("} ") + if not value: + continue + found = {} + for item in labels[len(metric) + 1 :].split(","): + name, raw = item.split("=", 1) + found[name] = raw.strip('"') + if all(found.get(name) == expected for name, expected in wanted.items()): + return float(value.split()[0]) + return 0.0 + + +def rpc_count(port, method, statuses=("success",)): + text = fetch_text(port, "/metrics") + return sum(metric_value(text, method, status) for status in statuses) + + +def fetch_text(port, path): + conn = http.client.HTTPConnection(HOST, port, timeout=2) + try: + conn.request("GET", path) + response = conn.getresponse() + body = response.read() + if response.status != 200: + raise RuntimeError(f"GET {path} on {port}: HTTP {response.status}") + return body.decode() + finally: + conn.close() + + +def put_requests(http_port, key, body, count): + path = "/v1/cache/" + urllib.parse.quote(key, safe="") + conn = http.client.HTTPConnection(HOST, http_port, timeout=30) + latencies = [] + try: + for _ in range(count): + start = time.perf_counter() + conn.request( + "POST", + path, + body=body, + headers={"Content-Type": "application/json", "Content-Length": str(len(body))}, + ) + response = conn.getresponse() + response_body = response.read() + elapsed = time.perf_counter() - start + if response.status != 200: + raise RuntimeError(f"PUT {key}: HTTP {response.status}: {response_body[:200]!r}") + result = json.loads(response_body) + if not result.get("success", False): + raise RuntimeError(f"PUT {key}: unsuccessful response: {result}") + latencies.append(elapsed) + finally: + conn.close() + return latencies + + +def start_cluster(binary, root): + for attempt in range(5): + ports = free_ports(9) + grpc_ports = ports[0:3] + http_ports = ports[3:6] + gossip_ports = ports[6:9] + seeds = ",".join(f"{HOST}:{port}" for port in gossip_ports) + attempt_root = root / f"attempt-{attempt}" + attempt_root.mkdir() + processes = [] + try: + for index in range(3): + node_root = attempt_root / f"node{index + 1}" + node_root.mkdir() + log = (attempt_root / f"node{index + 1}.log").open("w") + command = [ + binary, + "-cluster-enabled", + "-node-id", + f"node{index + 1}", + "-disk", + str(node_root), + "-listen-addr", + f"{HOST}:{grpc_ports[index]}", + "-listen-http", + f"{HOST}:{http_ports[index]}", + "-cluster-addr", + f"{HOST}:{gossip_ports[index]}", + "-seeds", + seeds, + ] + processes.append(subprocess.Popen(command, stdout=log, stderr=log)) + wait_for_http(http_ports[index], "/ready") + log.close() + topology = wait_for_http(http_ports[0], "/v1/topology", topology_ready, timeout=90) + return processes, grpc_ports, http_ports, topology + except Exception: + stop_cluster(processes) + if attempt == 4: + raise + time.sleep(0.2) + + +def stop_cluster(processes): + for process in processes: + if process.poll() is None: + process.terminate() + for process in processes: + try: + process.wait(timeout=10) + except subprocess.TimeoutExpired: + process.kill() + process.wait() + + +def emit(metric, value): + print(json.dumps({"metric": metric, "value": value}, separators=(",", ":"))) + + +def main(): + parser = argparse.ArgumentParser() + parser.add_argument("--binary", required=True) + parser.add_argument("--size", type=int, required=True) + parser.add_argument("--requests", type=int, default=30) + args = parser.parse_args() + if args.size <= 0 or args.requests <= 0: + raise ValueError("size and requests must be positive") + + processes = [] + with tempfile.TemporaryDirectory(prefix="ocache-http-gateway-bench-") as temporary: + try: + processes, _, http_ports, topology = start_cluster(args.binary, pathlib.Path(temporary)) + tokens = node_tokens(topology) + key = None + owner = None + for candidate in (f"http-gateway-{args.size}-{index}" for index in range(10000)): + candidate_owner = owner_for(candidate, tokens) + if candidate_owner != "node1": + key = candidate + owner = candidate_owner + break + if key is None: + raise RuntimeError("could not find a key owned by a node other than node1") + owner_index = int(owner.removeprefix("node")) - 1 + + body = json.dumps( + {"data": base64.b64encode(b"x" * args.size).decode()}, + separators=(",", ":"), + ).encode() + put_requests(http_ports[0], key, body, 10) + + before_gateway_put = rpc_count(http_ports[0], "PutObject") + before_owner_put_object = rpc_count(http_ports[owner_index], "PutObject") + before_owner_put_stream = rpc_count(http_ports[owner_index], "Put", ("success", "forwarded")) + before_cpu = sum(scalar_metric(fetch_text(port, "/metrics"), "process_cpu_seconds_total") for port in http_ports) + + latencies = put_requests(http_ports[0], key, body, args.requests) + + after_cpu = sum(scalar_metric(fetch_text(port, "/metrics"), "process_cpu_seconds_total") for port in http_ports) + gateway_put = rpc_count(http_ports[0], "PutObject") - before_gateway_put + owner_put_object = rpc_count(http_ports[owner_index], "PutObject") - before_owner_put_object + owner_put_stream = rpc_count(http_ports[owner_index], "Put", ("success", "forwarded")) - before_owner_put_stream + cpu_us_per_request = (after_cpu - before_cpu) * 1_000_000 / args.requests + ordered = sorted(latencies) + p50_us = ordered[len(ordered) // 2] * 1_000_000 + + emit("http-put-p50-us", p50_us) + emit("server-cpu-us-per-request", cpu_us_per_request) + emit("gateway-putobject-per-request", gateway_put / args.requests) + emit("owner-putobject-per-request", owner_put_object / args.requests) + emit("owner-putstream-per-request", owner_put_stream / args.requests) + finally: + stop_cluster(processes) + + +if __name__ == "__main__": + try: + main() + except Exception as error: + print(f"benchmark failed: {error}", file=sys.stderr) + raise diff --git a/server/service/gateway.go b/server/service/gateway.go new file mode 100644 index 00000000..272008e8 --- /dev/null +++ b/server/service/gateway.go @@ -0,0 +1,93 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !ocache_topology_benchmark + +package service + +import ( + "context" + "errors" + + "github.com/tigrisdata/ocache/coordinator" + "github.com/tigrisdata/ocache/coordinator/ring" + pb "github.com/tigrisdata/ocache/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/status" +) + +// gatewayUnaryForwardLimit keeps the HTTP gateway's direct owner call on the +// bounded unary path. Larger bodies continue through the local service, whose +// existing Operations.PutBytes path forwards them with the streaming RPC. +const gatewayUnaryForwardLimit = 4 * 1024 * 1024 + +type gatewayRouter interface { + Resolve(key string) (*ring.NodeInfo, error) + RouteToNode(nodeID string) (pb.CacheServiceClient, error) + GetLocalNodeID() string +} + +// gatewayCacheServiceClient keeps the generated gateway on its normal local +// client for every method except the eligible unary PutObject path. Embedding +// the generated client means new and streaming methods retain their existing +// transport and handler behavior. +type gatewayCacheServiceClient struct { + pb.CacheServiceClient + router gatewayRouter +} + +func newGatewayCacheServiceClient(local pb.CacheServiceClient, coord *coordinator.Coordinator) *gatewayCacheServiceClient { + var router gatewayRouter + if coord != nil { + router = coord + } + return &gatewayCacheServiceClient{ + CacheServiceClient: local, + router: router, + } +} + +func (c *gatewayCacheServiceClient) PutObject(ctx context.Context, req *pb.PutRequest, opts ...grpc.CallOption) (*pb.PutResponse, error) { + if c.router == nil || req == nil || req.Key == "" || len(req.Data) > gatewayUnaryForwardLimit { + return c.CacheServiceClient.PutObject(ctx, req, opts...) + } + // Let the local gRPC interceptor preserve its transport-level response for + // requests that have already exceeded the forwarding limit. + if err := coordinator.CheckHopCount(coordinator.ExtractRequestMetadata(ctx).HopCount); err != nil { + return c.CacheServiceClient.PutObject(ctx, req, opts...) + } + + owner, err := c.router.Resolve(req.Key) + if err != nil || owner == nil || owner.ID == "" || owner.ID == c.router.GetLocalNodeID() { + return c.CacheServiceClient.PutObject(ctx, req, opts...) + } + + forwardedCtx, err := coordinator.IncrementHopCount(ctx, c.router.GetLocalNodeID()) + if err != nil { + return gatewayPutError(err) + } + + remote, err := c.router.RouteToNode(owner.ID) + if err != nil { + return gatewayPutError(err) + } + + resp, err := remote.PutObject(forwardedCtx, req, opts...) + if err != nil { + return gatewayPutError(err) + } + return resp, nil +} + +func gatewayPutError(err error) (*pb.PutResponse, error) { + if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) { + return nil, err + } + code := status.Code(err) + if code == codes.Canceled || code == codes.DeadlineExceeded { + return nil, err + } + userErr := mapStorageErrorToGRPC(err) + return &pb.PutResponse{Success: false, Error: userErr.Error()}, nil +} diff --git a/server/service/gateway_test.go b/server/service/gateway_test.go new file mode 100644 index 00000000..8eaef56a --- /dev/null +++ b/server/service/gateway_test.go @@ -0,0 +1,319 @@ +// Copyright 2026 Tigris Data, Inc. +// SPDX-License-Identifier: Apache-2.0 + +//go:build !ocache_topology_benchmark + +package service + +import ( + "bytes" + "context" + "errors" + "strconv" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" + "github.com/tigrisdata/ocache/coordinator" + "github.com/tigrisdata/ocache/coordinator/ring" + pb "github.com/tigrisdata/ocache/proto" + "google.golang.org/grpc" + "google.golang.org/grpc/codes" + "google.golang.org/grpc/metadata" + "google.golang.org/grpc/status" +) + +type gatewayTestClient struct { + pb.CacheServiceClient + putObject func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) +} + +func (c *gatewayTestClient) PutObject(ctx context.Context, req *pb.PutRequest, opts ...grpc.CallOption) (*pb.PutResponse, error) { + return c.putObject(ctx, req, opts...) +} + +type gatewayTestRouter struct { + owner *ring.NodeInfo + resolveErr error + routeErr error + localID string + remote pb.CacheServiceClient + resolveCall int + routeCall int +} + +func (r *gatewayTestRouter) Resolve(string) (*ring.NodeInfo, error) { + r.resolveCall++ + return r.owner, r.resolveErr +} + +func (r *gatewayTestRouter) RouteToNode(string) (pb.CacheServiceClient, error) { + r.routeCall++ + return r.remote, r.routeErr +} + +func (r *gatewayTestRouter) GetLocalNodeID() string { + return r.localID +} + +func TestGatewayCacheServiceClient_RoutesEligibleRemotePutObject(t *testing.T) { + var localCalls, remoteCalls int + var gotCtx context.Context + var gotReq *pb.PutRequest + var gotOpts int + + local := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + localCalls++ + return &pb.PutResponse{Success: true}, nil + }} + remote := &gatewayTestClient{putObject: func(ctx context.Context, req *pb.PutRequest, opts ...grpc.CallOption) (*pb.PutResponse, error) { + remoteCalls++ + gotCtx = ctx + gotReq = req + gotOpts = len(opts) + return &pb.PutResponse{Success: true}, nil + }} + router := &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + remote: remote, + } + client := &gatewayCacheServiceClient{CacheServiceClient: local, router: router} + + ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs( + coordinator.MetadataKeyHop, "1", + coordinator.MetadataKeyOrigin, "origin", + "authorization", "bearer", + )) + req := &pb.PutRequest{Key: "remote-key", Data: []byte("value"), TtlSeconds: 7} + wantOption := grpc.StaticMethod() + resp, err := client.PutObject(ctx, req, wantOption) + + require.NoError(t, err) + require.NotNil(t, resp) + assert.True(t, resp.Success) + assert.Zero(t, localCalls) + assert.Equal(t, 1, remoteCalls) + assert.Equal(t, 1, router.resolveCall) + assert.Equal(t, 1, router.routeCall) + assert.Same(t, req, gotReq) + assert.Equal(t, 1, gotOpts) + + outgoing, ok := metadata.FromOutgoingContext(gotCtx) + require.True(t, ok) + assert.Equal(t, []string{"2"}, outgoing.Get(coordinator.MetadataKeyHop)) + assert.Equal(t, []string{"true"}, outgoing.Get(coordinator.MetadataKeyForwarded)) + assert.Equal(t, []string{"origin"}, outgoing.Get(coordinator.MetadataKeyOrigin)) + assert.Equal(t, []string{"bearer"}, outgoing.Get("authorization")) +} + +func TestGatewayCacheServiceClientKeepsOrdinaryPutObjectOnLocalClient(t *testing.T) { + tests := []struct { + name string + req *pb.PutRequest + router *gatewayTestRouter + wantResolve bool + }{ + { + name: "cluster disabled", + req: &pb.PutRequest{Key: "key", Data: []byte("value")}, + }, + { + name: "invalid key", + req: &pb.PutRequest{Data: []byte("value")}, + router: &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + }, + }, + { + name: "local owner", + req: &pb.PutRequest{Key: "key", Data: []byte("value")}, + wantResolve: true, + router: &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "gateway"}, + localID: "gateway", + }, + }, + { + name: "large body", + req: &pb.PutRequest{Key: "key", Data: bytes.Repeat([]byte("x"), gatewayUnaryForwardLimit+1)}, + router: &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + }, + }, + { + name: "owner lookup failure", + req: &pb.PutRequest{Key: "key", Data: []byte("value")}, + wantResolve: true, + router: &gatewayTestRouter{ + localID: "gateway", + resolveErr: errors.New("ring unavailable"), + }, + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + localCalls := 0 + local := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + localCalls++ + return &pb.PutResponse{Success: true}, nil + }} + var router gatewayRouter + if tt.router != nil { + router = tt.router + } + client := &gatewayCacheServiceClient{CacheServiceClient: local, router: router} + + resp, err := client.PutObject(context.Background(), tt.req) + + require.NoError(t, err) + assert.True(t, resp.Success) + assert.Equal(t, 1, localCalls) + if tt.router != nil { + assert.Zero(t, tt.router.routeCall) + if tt.wantResolve { + assert.Equal(t, 1, tt.router.resolveCall) + } else { + assert.Zero(t, tt.router.resolveCall) + } + } + }) + } +} + +func TestGatewayCacheServiceClientDelegatesOverHopLimit(t *testing.T) { + localCalls := 0 + remoteCalls := 0 + local := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + localCalls++ + return &pb.PutResponse{Success: true}, nil + }} + remote := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + remoteCalls++ + return &pb.PutResponse{Success: true}, nil + }} + router := &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + remote: remote, + } + client := &gatewayCacheServiceClient{CacheServiceClient: local, router: router} + ctx := metadata.NewOutgoingContext(context.Background(), metadata.Pairs( + coordinator.MetadataKeyHop, strconv.Itoa(coordinator.MaxHops+1), + )) + + resp, err := client.PutObject(ctx, &pb.PutRequest{Key: "key", Data: []byte("value")}) + + require.NoError(t, err) + assert.True(t, resp.Success) + assert.Equal(t, 1, localCalls) + assert.Zero(t, remoteCalls) + assert.Zero(t, router.resolveCall) +} + +func TestGatewayCacheServiceClientTranslatesForwardingFailures(t *testing.T) { + tests := []struct { + name string + routeErr error + remoteErr error + wantError string + }{ + { + name: "route failure", + routeErr: errors.New("owner unavailable"), + wantError: "owner unavailable", + }, + { + name: "remote failure", + remoteErr: errors.New("rpc failed"), + wantError: "rpc failed", + }, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + local := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + return &pb.PutResponse{Success: true}, nil + }} + remote := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + return nil, tt.remoteErr + }} + router := &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + remote: remote, + routeErr: tt.routeErr, + } + client := &gatewayCacheServiceClient{CacheServiceClient: local, router: router} + + resp, err := client.PutObject(context.Background(), &pb.PutRequest{Key: "key", Data: []byte("value")}) + + require.NoError(t, err) + require.NotNil(t, resp) + assert.False(t, resp.Success) + assert.Equal(t, tt.wantError, resp.Error) + }) + } +} + +func TestGatewayCacheServiceClientPreservesCancellationFailures(t *testing.T) { + tests := []struct { + name string + err error + code codes.Code + }{ + {name: "context canceled", err: context.Canceled, code: codes.Canceled}, + {name: "grpc deadline", err: status.Error(codes.DeadlineExceeded, "owner deadline"), code: codes.DeadlineExceeded}, + } + + for _, tt := range tests { + t.Run(tt.name, func(t *testing.T) { + local := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + return &pb.PutResponse{Success: true}, nil + }} + remote := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + return nil, tt.err + }} + router := &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + remote: remote, + } + client := &gatewayCacheServiceClient{CacheServiceClient: local, router: router} + + resp, err := client.PutObject(context.Background(), &pb.PutRequest{Key: "key", Data: []byte("value")}) + + assert.Nil(t, resp) + require.Error(t, err) + if errors.Is(tt.err, context.Canceled) { + assert.ErrorIs(t, err, context.Canceled) + } else { + assert.Equal(t, tt.code, status.Code(err)) + } + }) + } +} + +func TestGatewayCacheServiceClientPreservesRemoteResponse(t *testing.T) { + local := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + return &pb.PutResponse{Success: true}, nil + }} + remoteResponse := &pb.PutResponse{Success: false, Error: "owner rejected write"} + remote := &gatewayTestClient{putObject: func(context.Context, *pb.PutRequest, ...grpc.CallOption) (*pb.PutResponse, error) { + return remoteResponse, nil + }} + router := &gatewayTestRouter{ + owner: &ring.NodeInfo{ID: "owner"}, + localID: "gateway", + remote: remote, + } + client := &gatewayCacheServiceClient{CacheServiceClient: local, router: router} + + resp, err := client.PutObject(context.Background(), &pb.PutRequest{Key: "key", Data: []byte("value")}) + + require.NoError(t, err) + assert.Same(t, remoteResponse, resp) +} diff --git a/server/service/service.go b/server/service/service.go index cd5c3554..43b9dc03 100644 --- a/server/service/service.go +++ b/server/service/service.go @@ -207,11 +207,27 @@ func StartGRPCGatewayServer(coord *coordinator.Coordinator, grpcAddr string, lis mux := http.NewServeMux() gwMux := runtime.NewServeMux() - // Register the gRPC service handler with grpc-gateway + // Keep the generated gateway client for ordinary requests, but dispatch + // eligible nonlocal unary puts directly to their resolved owner. opts := []grpc.DialOption{grpc.WithInsecure()} - if err := pb.RegisterCacheServiceHandlerFromEndpoint(ctx, gwMux, grpcAddr, opts); err != nil { + localConn, err := grpc.NewClient(grpcAddr, opts...) + if err != nil { + zlog.Fatal().Err(err).Msg("failed to connect grpc-gateway handler") + } + localClient := pb.NewCacheServiceClient(localConn) + gatewayClient := newGatewayCacheServiceClient(localClient, coord) + if err := pb.RegisterCacheServiceHandlerClient(ctx, gwMux, gatewayClient); err != nil { + if closeErr := localConn.Close(); closeErr != nil { + zlog.Error().Err(closeErr).Msg("failed to close grpc-gateway client") + } zlog.Fatal().Err(err).Msg("failed to register grpc-gateway handler") } + go func() { + <-ctx.Done() + if err := localConn.Close(); err != nil { + zlog.Error().Err(err).Msg("failed to close grpc-gateway client") + } + }() // Health check endpoints for Docker/Kubernetes mux.HandleFunc("/health", healthHandler())