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)