From 6cfb78707442eee79ffcc34db7af98e68ebea9e4 Mon Sep 17 00:00:00 2001 From: "David Muto (pseudomuto)" Date: Wed, 15 Jul 2026 08:23:39 -0400 Subject: [PATCH] [router]: Dynamic per-request routing to upstream The router previously forwarded every stream to a single upstream resolved from PrimaryUpstream, so a multi-upstream config had no way to send a request to anything but the first upstream. The handler now routes each request on its own. It buffers the first client frame, uses the Reflector (backed by the proto Extractor) to peek the request namespace from that frame, and hands the namespace and metadata to a Director. The module's director matches them against the Mux compiled from the routing config and returns the pooled connection for the selected upstream, dialing one connection per configured upstream up front. Unroutable requests fail with FailedPrecondition, and a matched upstream with no connection fails with Unavailable. --- .github/CONTRIBUTING.md | 59 +++++++-- .mise/config.toml | 13 ++ dev/Procfile | 4 +- dev/config.yaml | 28 ++++- internal/config/config.go | 11 -- internal/config/config_test.go | 24 ---- internal/router/fx.go | 98 +++++++++++---- internal/router/fx_test.go | 4 + internal/router/handler.go | 69 ++++++++-- internal/router/handler_test.go | 217 +++++++++++++++++++++++++++++++- 10 files changed, 440 insertions(+), 87 deletions(-) diff --git a/.github/CONTRIBUTING.md b/.github/CONTRIBUTING.md index efd970c..fed578b 100644 --- a/.github/CONTRIBUTING.md +++ b/.github/CONTRIBUTING.md @@ -21,20 +21,27 @@ pinned dev tools (golangci-lint, mockgen). ## Testing the proxy locally -`mise run server` starts everything you need: a Temporal dev server (the upstream, on `localhost:7233` with namespaces -`ns1.remote` and `ns2.remote`) and the proxy itself, listening on `localhost:8444` with TLS terminated using the dev -certs in `dev/certs/`. The inbound server requires a client certificate (mTLS), so requests must present -`dev/certs/client.crt`. +`mise run server` starts everything you need: three Temporal dev servers and the proxy. Each dev server backs one +upstream defined in `dev/config.yaml`: + +| Dev server | Namespaces | Upstream | +| ---------------- | -------------------------- | ----------- | +| `localhost:7233` | `ns1.remote`, `ns2.remote` | `cluster-1` | +| `localhost:7234` | `test` | `cluster-2` | +| `localhost:7235` | `test2` | `cluster-3` | + +The proxy itself listens on `localhost:8444` with TLS terminated using the dev certs in `dev/certs/`. The inbound server +requires a client certificate (mTLS), so requests must present `dev/certs/client.crt`. In a second shell, use the `grpc` task to send requests. It wraps `buf curl` with the dev certs and gRPC-over-HTTP2 flags: ```sh -mise run grpc [json-body] +mise run grpc [json-body] [-H 'key: value' ...] ``` A bare method name is assumed to be on `WorkflowService`. Anything containing a slash is used as a full `service/method` -path. +path. Pass `-H`/`--header` (repeatable) to attach gRPC metadata, for example `-H 'x-cluster: 3'`. ### Confirming requests are forwarded upstream @@ -52,9 +59,43 @@ To exercise a namespace-scoped call, use `DescribeNamespace`: mise run grpc DescribeNamespace '{"namespace":"ns1.remote"}' ``` -> **Note:** namespace names are currently forwarded verbatim - the `upstream.namespaces.rules` in `dev/config.yaml` are -> parsed and validated but not yet applied. Until translation is wired up, you must send the real upstream name -> (`ns1.remote`), not the local alias (`ns1`). Sending `ns1` returns `NotFound`. +> [!IMPORTANT] +> +> namespace names are currently forwarded verbatim - the `upstream.namespaces.rules` in `dev/config.yaml` are parsed and +> validated but not yet applied. Until translation is wired up, you must send the real upstream name (`ns1.remote`), not +> the local alias (`ns1`). Sending `ns1` returns `NotFound`. + +### Confirming per-request routing + +The proxy picks an upstream per request from the `routing` rules in `dev/config.yaml`, evaluated in order (first match +wins): + +1. namespace `test` -> `cluster-2` +2. metadata `x-cluster: 3` -> `cluster-3` + +Everything else, including requests with no namespace (such as `GetSystemInfo`), falls through to the `default` upstream +(`cluster-1`). A namespace-scoped call to `test` lands on the second dev server: + +```sh +mise run grpc DescribeNamespace '{"namespace":"test"}' +``` + +#### Routing on metadata + +The second rule routes on request metadata rather than namespace, so any request carrying `x-cluster: 3` goes to +`cluster-3` regardless of namespace. `cluster-3` (`localhost:7235`) is the only dev server with the `test2` namespace, +which makes the routing observable: + +```sh +# Routed to cluster-3 (which has test2) -> succeeds +mise run grpc DescribeNamespace '{"namespace":"test2"}' -H 'x-cluster: 3' + +# No header -> falls through to cluster-1 (which does not have test2) -> NotFound +mise run grpc DescribeNamespace '{"namespace":"test2"}' +``` + +The differing results confirm the metadata rule selected the upstream. Metadata keys are matched case-insensitively +(gRPC lowercases them). ### Checking the front door (not proxied) diff --git a/.mise/config.toml b/.mise/config.toml index 2bbc834..0251c90 100644 --- a/.mise/config.toml +++ b/.mise/config.toml @@ -47,6 +47,7 @@ depends = ["gen-certs"] usage = ''' arg "" help="Method name (e.g. GetSystemInfo) or a full service/method path" arg "[data]" help="Request body as JSON" default="{}" +flag "-H --header
" help="Repeatable gRPC header, e.g. 'x-cluster: 3'" var=#true ''' run = ''' case "$usage_method" in @@ -54,11 +55,23 @@ case "$usage_method" in *) path="temporal.api.workflowservice.v1.WorkflowService/$usage_method" ;; esac +# mise exposes the repeatable --header flag as a shell-quoted list. Rebuild the +# positional parameters as interleaved "-H " pairs so header values that +# contain spaces (e.g. "x-cluster: 3") survive intact. +eval "set -- $usage_header" +n=$# +while [ "$n" -gt 0 ]; do + h="$1"; shift + set -- "$@" -H "$h" + n=$((n - 1)) +done + buf curl --protocol grpc --http2-prior-knowledge \ --cacert dev/certs/ca.crt \ --cert dev/certs/client.crt \ --key dev/certs/client.key \ --schema buf.build/temporalio/api \ + "$@" \ -d "$usage_data" \ "https://localhost:8444/$path" ''' diff --git a/dev/Procfile b/dev/Procfile index e3f3970..6007908 100644 --- a/dev/Procfile +++ b/dev/Procfile @@ -1,2 +1,4 @@ -temporal: temporal server start-dev -n ns1.remote -n ns2.remote +temporal1: temporal server start-dev -n ns1.remote -n ns2.remote +temporal2: temporal server start-dev -p 7234 -n test +temporal3: temporal server start-dev -p 7235 -n test2 proxy: go run ./cmd/proxy serve -c dev/config.yaml diff --git a/dev/config.yaml b/dev/config.yaml index e46af5c..ad4226f 100644 --- a/dev/config.yaml +++ b/dev/config.yaml @@ -5,7 +5,7 @@ tls: key: dev/certs/server.key upstreams: - - name: default + - name: cluster-1 hostPort: localhost:7233 namespaces: rules: @@ -16,5 +16,29 @@ upstreams: - local: ns3 remote: ns2.remote + - name: cluster-2 + hostPort: localhost:7234 + namespaces: + rules: + # Always add .remote on the way out, and remove it on the way in. + suffix: .remote + overrides: + # Makes ns2 and ns3 both use ns2.remote on the remote cluster. + - local: ns3 + remote: ns2.remote + + - name: cluster-3 + hostPort: localhost:7235 + routing: - default: default + default: cluster-1 + + rules: + - upstream: cluster-2 + match: + namespace: test + + - upstream: cluster-3 + match: + metadata: + x-cluster: "3" diff --git a/internal/config/config.go b/internal/config/config.go index 5d0e310..0de414b 100644 --- a/internal/config/config.go +++ b/internal/config/config.go @@ -84,14 +84,3 @@ func (c *Config) Validate() error { rules = append(rules, c.Routing.referentialRules(known)...) return validation.Validate("", rules...) } - -// PrimaryUpstream returns the first configured upstream. It is a temporary -// bridge for the single-upstream wiring in the proxy and router modules until -// per-request routing is in place. -func (c *Config) PrimaryUpstream() (*Upstream, error) { - if len(c.Upstreams) == 0 { - return nil, errors.New("no upstreams configured") - } - - return &c.Upstreams[0], nil -} diff --git a/internal/config/config_test.go b/internal/config/config_test.go index ca7b20a..59527a9 100644 --- a/internal/config/config_test.go +++ b/internal/config/config_test.go @@ -338,30 +338,6 @@ func TestConfig_Validate_RoutingReferences(t *testing.T) { } } -func TestConfig_PrimaryUpstream(t *testing.T) { - t.Parallel() - - t.Run("returns the first configured upstream", func(t *testing.T) { - t.Parallel() - - cfg := &config.Config{Upstreams: []config.Upstream{ - {Name: "first", Listen: config.ListenConfig{HostPort: "127.0.0.1:7233"}}, - {Name: "second", Listen: config.ListenConfig{HostPort: "127.0.0.1:7234"}}, - }} - - up, err := cfg.PrimaryUpstream() - require.NoError(t, err) - require.Equal(t, "first", up.Name) - }) - - t.Run("errors when no upstreams are configured", func(t *testing.T) { - t.Parallel() - - _, err := (&config.Config{}).PrimaryUpstream() - require.Error(t, err) - }) -} - func TestConfig_ValidateRejectsDuplicateHostPorts(t *testing.T) { t.Parallel() diff --git a/internal/router/fx.go b/internal/router/fx.go index 2773082..f41aa47 100644 --- a/internal/router/fx.go +++ b/internal/router/fx.go @@ -1,46 +1,59 @@ package router import ( + "context" "fmt" "strings" "go.uber.org/fx" "google.golang.org/grpc" + "google.golang.org/grpc/codes" "google.golang.org/grpc/credentials/insecure" + "google.golang.org/grpc/status" "github.com/temporalio/temporal-proxy/internal/config" + "github.com/temporalio/temporal-proxy/internal/protoutil" "github.com/temporalio/temporal-proxy/internal/transport/connect" "github.com/temporalio/temporal-proxy/internal/transport/socket" "github.com/temporalio/temporal-proxy/pkg/match" ) -// Module is the fx module that provides the transparent-forwarding pieces: a -// pass-through [google.golang.org/grpc/encoding.CodecV2] and a -// [google.golang.org/grpc.StreamHandler]. The handler obtains a connection to -// the proxy's unix socket, whose path is derived from the upstream host:port in -// configuration, from the shared [connect.Pool]. +// Module is the fx module that provides the routing-and-forwarding pieces: a +// pass-through [google.golang.org/grpc/encoding.CodecV2], a [Mux] compiled from +// the routing configuration, and a [google.golang.org/grpc.StreamHandler]. The +// handler dials one connection per configured upstream from the shared +// [connect.Pool] (each unix socket path derived from that upstream's +// host:port), then routes every request to an upstream by matching it with the +// Mux. var Module = fx.Options(fx.Provide( Codec, func(p RouterParams) (grpc.StreamHandler, error) { - upstream, err := p.Config.PrimaryUpstream() - if err != nil { - return nil, fmt.Errorf("failed to resolve upstream: %w", err) - } + conns := make(map[string]*grpc.ClientConn, len(p.Config.Upstreams)) + for i := range p.Config.Upstreams { + upstream := &p.Config.Upstreams[i] + sockPath, err := socket.UnixPath(upstream.Listen.HostPort) + if err != nil { + return nil, fmt.Errorf("failed to resolve proxy socket path[%q]: %w", upstream.Name, err) + } - sockPath, err := socket.UnixPath(upstream.Listen.HostPort) - if err != nil { - return nil, fmt.Errorf("failed to resolve proxy socket path: %w", err) - } + conn, err := p.Pool.GetOrSet( + "unix://"+sockPath, + grpc.WithTransportCredentials(insecure.NewCredentials()), + ) + if err != nil { + return nil, fmt.Errorf("failed to create upstream client[%q]: %w", upstream.Name, err) + } - conn, err := p.Pool.GetOrSet( - "unix://"+sockPath, - grpc.WithTransportCredentials(insecure.NewCredentials()), - ) - if err != nil { - return nil, fmt.Errorf("failed to create upstream client: %w", err) + conns[upstream.Name] = conn } - return Handler(conn), nil + return Handler( + &director{ + conns: conns, + mux: p.Mux, + }, + p.Extractor, + ), nil }, func(c *config.Config) (*Mux, error) { rules := make([]Rule, 0, len(c.Routing.Rules)) @@ -90,11 +103,44 @@ var Module = fx.Options(fx.Provide( }, )) -// RouterParams collects the fx-provided dependencies needed to build the -// forwarding stream handler. -type RouterParams struct { - fx.In +type ( + // RouterParams collects the fx-provided dependencies needed to build the + // forwarding stream handler. + RouterParams struct { + fx.In + + Config *config.Config + Extractor *protoutil.Extractor + Mux *Mux + Pool *connect.Pool + } + + // director is the [Director] used by the module's handler. It maps the + // upstream name chosen by the Mux to that upstream's pooled connection. + director struct { + conns map[string]*grpc.ClientConn + mux *Mux + } +) + +// Resolve routes a request by matching it against the Mux and returning the +// connection for the resulting upstream. It fails with FailedPrecondition when +// no upstream matches (and no default is configured) and with Unavailable when +// the matched upstream has no connection. +func (d *director) Resolve( + ctx context.Context, + _, namespace string, + md map[string][]string, +) (*grpc.ClientConn, error) { + upstream := d.mux.Switch(namespace, md) + if upstream == "" { + return nil, status.Error(codes.FailedPrecondition, "no upstream matched the request and no default is configured") + } + + cc, ok := d.conns[upstream] + if !ok { + return nil, status.Errorf(codes.Unavailable, "router: no connection for upstream %q", upstream) + } - Config *config.Config - Pool *connect.Pool + return cc, nil } diff --git a/internal/router/fx_test.go b/internal/router/fx_test.go index e3357c9..8be6637 100644 --- a/internal/router/fx_test.go +++ b/internal/router/fx_test.go @@ -15,6 +15,7 @@ import ( "google.golang.org/grpc/test/bufconn" "github.com/temporalio/temporal-proxy/internal/config" + "github.com/temporalio/temporal-proxy/internal/protoutil" "github.com/temporalio/temporal-proxy/internal/router" "github.com/temporalio/temporal-proxy/internal/transport/connect" "github.com/temporalio/temporal-proxy/internal/transport/socket" @@ -36,6 +37,7 @@ func TestModule(t *testing.T) { Upstreams: []config.Upstream{{Name: "primary", Listen: config.ListenConfig{HostPort: "127.0.0.1:7233"}}}, }), connect.Module, + protoutil.Module, router.Module, fx.Populate(&codec, &handler), fx.NopLogger, @@ -98,9 +100,11 @@ func TestModule(t *testing.T) { ) app := fx.New( fx.Supply(&config.Config{ + Routing: config.Routing{DefaultUpstream: "primary"}, Upstreams: []config.Upstream{{Name: "primary", Listen: config.ListenConfig{HostPort: upstream}}}, }), connect.Module, + protoutil.Module, router.Module, fx.Populate(&codec, &handler), fx.NopLogger, diff --git a/internal/router/handler.go b/internal/router/handler.go index 89e3428..5304461 100644 --- a/internal/router/handler.go +++ b/internal/router/handler.go @@ -1,7 +1,9 @@ package router import ( + "context" "io" + "maps" "google.golang.org/grpc" "google.golang.org/grpc/codes" @@ -9,26 +11,65 @@ import ( "google.golang.org/grpc/status" ) +type ( + // Director selects the upstream connection for a request. Resolve receives + // the full method, the namespace peeked from the first request message + // (empty when the client sent no message), and the incoming metadata, and + // returns the connection to forward the stream over. A non-nil error aborts + // the stream and is returned to the caller verbatim, so implementations + // should return a gRPC status error. + Director interface { + Resolve(ctx context.Context, method, namespace string, md map[string][]string) (*grpc.ClientConn, error) + } + + // Reflector extracts the Temporal namespace from a request. Namespace + // receives the full method and the raw bytes of the first request message + // and returns the namespace, or "" when it cannot determine one. + Reflector interface { + Namespace(string, []byte) string + } +) + // Handler returns a grpc.StreamHandler suitable for grpc.UnknownServiceHandler. -// It transparently forwards every stream to cc using the same full method name, -// pumping raw frames in both directions and propagating header, trailer, and +// It buffers the first request frame so r can peek the request namespace, asks d +// for the upstream connection, then transparently forwards the stream to that +// upstream using the same full method name: it replays the buffered first frame, +// pumps raw frames in both directions, and propagates header, trailer, and // status verbatim. -// -// cc is resolved once per stream, which is the seam where future per-request -// routing (selecting a connection from request details) will plug in. -func Handler(cc *grpc.ClientConn) grpc.StreamHandler { +func Handler(d Director, r Reflector) grpc.StreamHandler { return func(_ any, serverStream grpc.ServerStream) error { ctx := serverStream.Context() - sts := grpc.ServerTransportStreamFromContext(ctx) if sts == nil { return status.Error(codes.Internal, "router: no server transport stream in context") } - method := sts.Method() + var md map[string][]string + method := sts.Method() outCtx := ctx - if md, ok := metadata.FromIncomingContext(ctx); ok { - outCtx = metadata.NewOutgoingContext(ctx, md.Copy()) + if inMD, ok := metadata.FromIncomingContext(ctx); ok { + outCtx = metadata.NewOutgoingContext(ctx, inMD.Copy()) + md = inMD + } + + // Buffer the first client frame so we can read the namespace before + // choosing an upstream. io.EOF means the client half-closed without + // sending a message (namespace is empty). + first := &frame{} + firstErr := serverStream.RecvMsg(first) + eof := firstErr == io.EOF + if firstErr != nil && !eof { + return StatusError(firstErr) + } + + namespace := "" + if !eof { + namespace = r.Namespace(method, first.payload) + } + + cc, err := d.Resolve(ctx, method, namespace, maps.Clone(md)) + if err != nil { + return err } stream, err := cc.NewStream( @@ -41,6 +82,14 @@ func Handler(cc *grpc.ClientConn) grpc.StreamHandler { return err } + if eof { + if err := stream.CloseSend(); err != nil { + return StatusError(err) + } + } else if err := stream.SendMsg(first); err != nil { + return StatusError(err) + } + reqErr := pumpServerToClient(serverStream, stream) respErr := pumpClientToServer(stream, serverStream) diff --git a/internal/router/handler_test.go b/internal/router/handler_test.go index f4c9afd..4bdd7f3 100644 --- a/internal/router/handler_test.go +++ b/internal/router/handler_test.go @@ -1,10 +1,12 @@ package router_test import ( + "bytes" "context" "errors" "io" "net" + "sync" "testing" "github.com/stretchr/testify/require" @@ -16,11 +18,42 @@ import ( "google.golang.org/grpc/metadata" "google.golang.org/grpc/status" "google.golang.org/grpc/test/bufconn" + "google.golang.org/protobuf/proto" "github.com/temporalio/temporal-proxy/internal/router" "github.com/temporalio/temporal-proxy/internal/server" ) +type ( + // recordingDirector captures what Resolve was called with and returns a fixed + // connection. + recordingDirector struct { + cc *grpc.ClientConn + + mu sync.Mutex + calls int + method string + namespace string + md map[string][]string + } + + // recordingReflector captures what Namespace was called with and returns a + // fixed namespace. It is safe for the handler goroutine to write while the test + // goroutine reads via snapshot. + recordingReflector struct { + ns string + + mu sync.Mutex + calls int + method string + payload []byte + } + + stubDirector struct{ cc *grpc.ClientConn } + + stubReflector struct{} +) + func TestHandlerForwardsUnary(t *testing.T) { t.Parallel() @@ -66,6 +99,127 @@ func TestHandlerPropagatesError(t *testing.T) { require.Equal(t, codes.NotFound, status.Code(err)) } +func TestHandlerRoutesUsingReflectorAndDirector(t *testing.T) { + t.Parallel() + + echoDesc := grpc.ServiceDesc{ + ServiceName: "test.v1.Echo", + HandlerType: (*any)(nil), + Methods: []grpc.MethodDesc{ + { + MethodName: "Ping", + Handler: func(_ any, _ context.Context, dec func(any) error, _ grpc.UnaryServerInterceptor) (any, error) { + in := new(grpc_health_v1.HealthCheckRequest) + if err := dec(in); err != nil { + return nil, err + } + return &grpc_health_v1.HealthCheckResponse{Status: grpc_health_v1.HealthCheckResponse_SERVING}, nil + }, + }, + }, + } + + reflector := &recordingReflector{ns: "ns-from-reflector"} + var director *recordingDirector + relay := newRelayWith( + t, + func(s *grpc.Server) { s.RegisterService(&echoDesc, nil) }, + func(cc *grpc.ClientConn) router.Director { + director = &recordingDirector{cc: cc} + return director + }, + reflector, + ) + + ctx := metadata.AppendToOutgoingContext(t.Context(), "x-route", "gold") + resp := new(grpc_health_v1.HealthCheckResponse) + require.NoError(t, relay.Invoke(ctx, "/test.v1.Echo/Ping", &grpc_health_v1.HealthCheckRequest{Service: "abc"}, resp)) + require.Equal(t, grpc_health_v1.HealthCheckResponse_SERVING, resp.GetStatus()) + + // The Reflector is handed the method and the raw bytes of the first frame. + rCalls, rMethod, rPayload := reflector.snapshot() + require.Equal(t, 1, rCalls) + require.Equal(t, "/test.v1.Echo/Ping", rMethod) + gotReq := new(grpc_health_v1.HealthCheckRequest) + require.NoError(t, proto.Unmarshal(rPayload, gotReq)) + require.Equal(t, "abc", gotReq.GetService()) + + // The Director is handed the extracted namespace and the incoming metadata. + dCalls, dMethod, dNS, dMD := director.snapshot() + require.Equal(t, 1, dCalls) + require.Equal(t, "/test.v1.Echo/Ping", dMethod) + require.Equal(t, "ns-from-reflector", dNS) + require.Equal(t, []string{"gold"}, dMD["x-route"]) +} + +func TestHandlerForwardsEmptyMessageHalfClose(t *testing.T) { + t.Parallel() + + // Sum reads request messages until the client half-closes, then reports how + // many it saw. It lets the test prove the upstream observed the half-close + // (it returns rather than blocking) and received no injected first frame. + countDesc := grpc.ServiceDesc{ + ServiceName: "test.v1.Count", + HandlerType: (*any)(nil), + Streams: []grpc.StreamDesc{ + { + StreamName: "Sum", + ClientStreams: true, + Handler: func(_ any, stream grpc.ServerStream) error { + n := 0 + for { + err := stream.RecvMsg(new(grpc_health_v1.HealthCheckRequest)) + if err == io.EOF { + break + } + if err != nil { + return err + } + n++ + } + return stream.SendMsg(&grpc_health_v1.HealthCheckResponse{ + Status: grpc_health_v1.HealthCheckResponse_ServingStatus(n), + }) + }, + }, + }, + } + + reflector := &recordingReflector{ns: "unused"} + var director *recordingDirector + relay := newRelayWith( + t, + func(s *grpc.Server) { s.RegisterService(&countDesc, nil) }, + func(cc *grpc.ClientConn) router.Director { + director = &recordingDirector{cc: cc} + return director + }, + reflector, + ) + + stream, err := relay.NewStream( + t.Context(), + &grpc.StreamDesc{ClientStreams: true, ServerStreams: true}, + "/test.v1.Count/Sum", + ) + require.NoError(t, err) + require.NoError(t, stream.CloseSend()) // Half-close without ever sending a message. + + resp := new(grpc_health_v1.HealthCheckResponse) + require.NoError(t, stream.RecvMsg(resp)) + require.Equal(t, 0, int(resp.GetStatus()), "upstream should observe zero request messages") + require.ErrorIs(t, stream.RecvMsg(new(grpc_health_v1.HealthCheckResponse)), io.EOF) + + // With no first frame, there is nothing to peek, so the Reflector is skipped + // and the Director routes with an empty namespace. + rCalls, _, _ := reflector.snapshot() + require.Zero(t, rCalls) + + dCalls, _, dNS, _ := director.snapshot() + require.Equal(t, 1, dCalls) + require.Empty(t, dNS) +} + func TestHandlerPropagatesHeaderAndTrailer(t *testing.T) { t.Parallel() @@ -174,7 +328,7 @@ func TestHandlerCoHostsLocalHealthWithForwarding(t *testing.T) { t.Cleanup(func() { _ = upstreamConn.Close() }) svr, err := server.New( - server.WithUnknownServiceHandler(router.Handler(upstreamConn)), + server.WithUnknownServiceHandler(router.Handler(stubDirector{upstreamConn}, stubReflector{})), server.WithServerCodec(router.Codec()), ) require.NoError(t, err) @@ -227,6 +381,43 @@ func TestStatusError(t *testing.T) { }) } +func (s stubDirector) Resolve(context.Context, string, string, map[string][]string) (*grpc.ClientConn, error) { + return s.cc, nil +} + +func (stubReflector) Namespace(string, []byte) string { return "" } + +func (r *recordingReflector) Namespace(method string, payload []byte) string { + r.mu.Lock() + defer r.mu.Unlock() + r.calls++ + r.method = method + r.payload = bytes.Clone(payload) + return r.ns +} + +func (r *recordingReflector) snapshot() (calls int, method string, payload []byte) { + r.mu.Lock() + defer r.mu.Unlock() + return r.calls, r.method, bytes.Clone(r.payload) +} + +func (d *recordingDirector) Resolve(_ context.Context, method, namespace string, md map[string][]string) (*grpc.ClientConn, error) { + d.mu.Lock() + defer d.mu.Unlock() + d.calls++ + d.method = method + d.namespace = namespace + d.md = md + return d.cc, nil +} + +func (d *recordingDirector) snapshot() (calls int, method, namespace string, md map[string][]string) { + d.mu.Lock() + defer d.mu.Unlock() + return d.calls, d.method, d.namespace, d.md +} + func dialBufconn(t *testing.T, lis *bufconn.Listener) *grpc.ClientConn { t.Helper() conn, err := grpc.NewClient( @@ -248,11 +439,29 @@ func serve(t *testing.T, srv *grpc.Server, lis *bufconn.Listener) { } // newRelayToUpstream stands up a fake upstream (configured by registerUpstream), -// then a bare relay server that forwards all methods to it via router.Handler. -// Returns a client conn pointed at the relay. +// then a bare relay server that forwards all methods to it via router.Handler +// using pass-through routing. Returns a client conn pointed at the relay. func newRelayToUpstream(t *testing.T, registerUpstream func(*grpc.Server)) *grpc.ClientConn { t.Helper() + return newRelayWith( + t, registerUpstream, + func(cc *grpc.ClientConn) router.Director { return stubDirector{cc} }, + stubReflector{}, + ) +} + +// newRelayWith is like newRelayToUpstream but lets the caller supply the Director +// and Reflector, so tests can observe how the handler routes. makeDirector +// receives the connection to the fake upstream. +func newRelayWith( + t *testing.T, + registerUpstream func(*grpc.Server), + makeDirector func(cc *grpc.ClientConn) router.Director, + reflector router.Reflector, +) *grpc.ClientConn { + t.Helper() + upstreamLis := bufconn.Listen(1024 * 1024) upstream := grpc.NewServer() registerUpstream(upstream) @@ -264,7 +473,7 @@ func newRelayToUpstream(t *testing.T, registerUpstream func(*grpc.Server)) *grpc relayLis := bufconn.Listen(1024 * 1024) relay := grpc.NewServer( grpc.ForceServerCodecV2(router.Codec()), - grpc.UnknownServiceHandler(router.Handler(upstreamConn)), + grpc.UnknownServiceHandler(router.Handler(makeDirector(upstreamConn), reflector)), ) serve(t, relay, relayLis)