[router]: Dynamic per-request routing to upstream

.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 <Method> [json-body]
+mise run grpc <Method> [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)
 

.mise/config.toml

@@ -47,18 +47,31 @@ depends = ["gen-certs"]
 usage = '''
 arg "<method>" help="Method name (e.g. GetSystemInfo) or a full service/method path"
 arg "[data]" help="Request body as JSON" default="{}"
+flag "-H --header <header>" help="Repeatable gRPC header, e.g. 'x-cluster: 3'" var=#true
 '''
 run = '''
 case "$usage_method" in
   */*) path="$usage_method" ;;
   *)   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 <value>" 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"
 '''

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

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"

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
-}

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()
 

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
 }

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,