arbitrary shard counts

.github/workflows/pull-request.yml

@@ -59,15 +59,15 @@ jobs:
           cache: ${{ github.ref == 'refs/heads/main' }} # only update the cache in main.
 
       - name: Run go build
-        run: go build ./...
+        run: make bins
 
       - name: Run go unittest
         run: make test
 
       - name: Install helm
         uses: azure/[email protected]
         with:
-          version: v3.17.3
+          version: v3.19.4
 
       - name: Install helm-unittest
         run: helm plugin install https://github.com/helm-unittest/helm-unittest.git

Makefile

@@ -11,7 +11,7 @@ GO_GET_TOOL    = go get -tool -modfile=$(TOOLS_MOD_FILE)
 
 # Disable cgo by default.
 CGO_ENABLED ?= 0
-TEST_ARG ?= -race -timeout=5m -tags test_dep
+TEST_ARG ?= -race -timeout=15m -tags test_dep -count=1
 BENCH_ARG ?= -benchtime=5000x
 
 ALL_SRC         := $(shell find . -name "*.go")

cmd/proxy/main.go

@@ -63,12 +63,17 @@ func buildCLIOptions() *cli.App {
 	return app
 }
 
-func startPProfHTTPServer(logger log.Logger, c config.ProfilingConfig) {
+func startPProfHTTPServer(logger log.Logger, c config.ProfilingConfig, proxyInstance *proxy.Proxy) {
 	addr := c.PProfHTTPAddress
 	if len(addr) == 0 {
 		return
 	}
 
+	// Add debug endpoint handler
+	http.HandleFunc("/debug/connections", func(w http.ResponseWriter, r *http.Request) {
+		proxy.HandleDebugInfo(w, r, proxyInstance, logger)
+	})
+
 	go func() {
 		logger.Info("Start pprof http server", tag.NewStringTag("address", addr))
 		if err := http.ListenAndServe(addr, nil); err != nil {
@@ -101,7 +106,7 @@ func startProxy(c *cli.Context) error {
 	}
 
 	cfg := proxyParams.ConfigProvider.GetS2SProxyConfig()
-	startPProfHTTPServer(proxyParams.Logger, cfg.ProfilingConfig)
+	startPProfHTTPServer(proxyParams.Logger, cfg.ProfilingConfig, proxyParams.Proxy)
 
 	if err := proxyParams.Proxy.Start(); err != nil {
 		return err

common/intra_headers.go

@@ -0,0 +1,37 @@
+package common
+
+import (
+	"context"
+
+	"google.golang.org/grpc/metadata"
+)
+
+const (
+	// Intra-proxy identification and tracing headers
+	IntraProxyHeaderKey           = "x-s2s-intra-proxy"
+	IntraProxyHeaderValue         = "1"
+	IntraProxyOriginProxyIDHeader = "x-s2s-origin-proxy-id"
+	IntraProxyHopCountHeader      = "x-s2s-hop-count"
+	IntraProxyTraceIDHeader       = "x-s2s-trace-id"
+)
+
+// IsIntraProxy checks incoming context metadata for intra-proxy marker.
+func IsIntraProxy(ctx context.Context) bool {
+	if md, ok := metadata.FromIncomingContext(ctx); ok {
+		if vals := md.Get(IntraProxyHeaderKey); len(vals) > 0 && vals[0] == IntraProxyHeaderValue {
+			return true
+		}
+	}
+	return false
+}
+
+// WithIntraProxyHeaders returns a new outgoing context with intra-proxy headers set.
+func WithIntraProxyHeaders(ctx context.Context, headers map[string]string) context.Context {
+	md, _ := metadata.FromOutgoingContext(ctx)
+	md = md.Copy()
+	md.Set(IntraProxyHeaderKey, IntraProxyHeaderValue)
+	for k, v := range headers {
+		md.Set(k, v)
+	}
+	return metadata.NewOutgoingContext(ctx, md)
+}

config/cluster_conn_config.go

@@ -16,6 +16,7 @@ type (
 		OutboundHealthCheck        HealthCheckConfig   `yaml:"outboundHealthCheck"`
 		InboundHealthCheck         HealthCheckConfig   `yaml:"inboundHealthCheck"`
 		ShardCountConfig           ShardCountConfig    `yaml:"shardCount"`
+		MemberlistConfig           *MemberlistConfig   `yaml:"memberlist"`
 	}
 	StringTranslator struct {
 		Mappings    []StringMapping `yaml:"mappings"`

config/config.go

@@ -40,6 +40,7 @@ type ShardCountMode string
 const (
 	ShardCountDefault ShardCountMode = ""
 	ShardCountLCM     ShardCountMode = "lcm"
+	ShardCountRouting ShardCountMode = "routing"
 )
 
 type HealthCheckProtocol string
@@ -153,7 +154,11 @@ type (
 		// TODO: Soon to be deprecated! Create an item in ClusterConnections instead
 		HealthCheck *HealthCheckConfig `yaml:"healthCheck"`
 		// TODO: Soon to be deprecated! Create an item in ClusterConnections instead
-		OutboundHealthCheck        *HealthCheckConfig    `yaml:"outboundHealthCheck"`
+		OutboundHealthCheck *HealthCheckConfig `yaml:"outboundHealthCheck"`
+		// TODO: Soon to be deprecated! Create an item in ClusterConnections instead
+		ShardCountConfig ShardCountConfig `yaml:"shardCount"`
+		// TODO: Soon to be deprecated! Create an item in ClusterConnections instead
+		MemberlistConfig           *MemberlistConfig     `yaml:"memberlist"`
 		NamespaceNameTranslation   NameTranslationConfig `yaml:"namespaceNameTranslation"`
 		SearchAttributeTranslation SATranslationConfig   `yaml:"searchAttributeTranslation"`
 		Metrics                    *MetricsConfig        `yaml:"metrics"`
@@ -217,6 +222,31 @@ type (
 	LoggingConfig struct {
 		ThrottleMaxRPS float64 `yaml:"throttleMaxRPS"`
 	}
+
+	MemberlistConfig struct {
+		// Enable distributed shard management using memberlist
+		Enabled bool `yaml:"enabled"`
+		// Node name for this proxy instance in the cluster
+		NodeName string `yaml:"nodeName"`
+		// Bind address for memberlist cluster communication
+		BindAddr string `yaml:"bindAddr"`
+		// Bind port for memberlist cluster communication
+		BindPort int `yaml:"bindPort"`
+		// List of existing cluster members to join
+		JoinAddrs []string `yaml:"joinAddrs"`
+		// Shard assignment strategy (deprecated - now uses actual ownership tracking)
+		ShardStrategy string `yaml:"shardStrategy"`
+		// Map of node names to their proxy service addresses for forwarding
+		ProxyAddresses map[string]string `yaml:"proxyAddresses"`
+		// Use TCP-only transport (disables UDP) for restricted networks
+		TCPOnly bool `yaml:"tcpOnly"`
+		// Disable TCP pings when using TCP-only mode
+		DisableTCPPings bool `yaml:"disableTCPPings"`
+		// Probe timeout for memberlist health checks
+		ProbeTimeoutMs int `yaml:"probeTimeoutMs"`
+		// Probe interval for memberlist health checks
+		ProbeIntervalMs int `yaml:"probeIntervalMs"`
+	}
 )
 
 func FromServerTLSConfig(cfg ServerTLSConfig) encryption.TLSConfig {

config/converter.go

@@ -41,6 +41,8 @@ func ToClusterConnConfig(config S2SProxyConfig) S2SProxyConfig {
 				SearchAttributeTranslation: config.SearchAttributeTranslation,
 				OutboundHealthCheck:        flattenNilHealthCheck(config.OutboundHealthCheck),
 				InboundHealthCheck:         flattenNilHealthCheck(config.HealthCheck),
+				ShardCountConfig:           config.ShardCountConfig,
+				MemberlistConfig:           config.MemberlistConfig,
 			},
 		},
 		Metrics:         config.Metrics,

develop/config/cluster-a-mux-client-proxy-1.yaml

@@ -0,0 +1,23 @@
+inbound:
+  name: "a-inbound-server"
+  server:
+    type: "mux"
+    mux: "muxed"
+  client:
+    tcp:
+      serverAddress: "localhost:7233"
+outbound:
+  name: "a-outbound-server"
+  server:
+    tcp:
+      listenAddress: "0.0.0.0:6133"
+  client:
+    type: "mux"
+    mux: "muxed"  
+mux:
+  - name: "muxed"
+    mode: "client"
+    client:
+      serverAddress: "localhost:7003" 
+profiling:
+  pprofAddress: "localhost:6060"

develop/config/cluster-a-mux-client-proxy-2.yaml

@@ -0,0 +1,23 @@
+inbound:
+  name: "a-inbound-server"
+  server:
+    type: "mux"
+    mux: "muxed"
+  client:
+    tcp:
+      serverAddress: "localhost:7233"
+outbound:
+  name: "a-outbound-server"
+  server:
+    tcp:
+      listenAddress: "0.0.0.0:6233"
+  client:
+    type: "mux"
+    mux: "muxed"  
+mux:
+  - name: "muxed"
+    mode: "client"
+    client:
+      serverAddress: "localhost:7003" 
+profiling:
+  pprofAddress: "localhost:6061"

develop/config/cluster-b-mux-server-proxy-1.yaml

@@ -0,0 +1,46 @@
+inbound:
+  name: "b-inbound-server"
+  server:
+    type: "mux"
+    mux: "muxed"
+  client:
+    tcp:
+      serverAddress: "localhost:8233"
+outbound:
+  name: "b-outbound-server"
+  server:
+    tcp:
+      listenAddress: "0.0.0.0:6333"
+  client:
+    type: "mux"
+    mux: "muxed"
+mux:
+  - name: "muxed"
+    mode: "server"
+    server:
+      listenAddress: "0.0.0.0:6334"
+# shardCount:
+#   mode: "lcm"
+#   localShardCount: 3
+#   remoteShardCount: 2
+shardCount:
+  mode: "routing"
+  localShardCount: 3
+  remoteShardCount: 2
+profiling:
+  pprofAddress: "localhost:6070"
+memberlist:
+  enabled: true
+  nodeName: "proxy-node-b-1"
+  bindAddr: "127.0.0.1"
+  bindPort: 6335
+  # joinAddrs:
+  #   - "localhost:6435"
+  proxyAddresses:
+    "proxy-node-b-1": "localhost:6333"
+    "proxy-node-b-2": "localhost:6433"
+  # # TCP-only configuration for restricted networks
+  tcpOnly: true           # Use TCP transport only, disable UDP
+  # disableTCPPings: true   # Disable TCP pings for faster convergence
+  # probeTimeoutMs: 1000    # Longer timeout for network latency
+  # probeIntervalMs: 2000   # Less frequent probes to reduce network noise

develop/config/cluster-b-mux-server-proxy-2.yaml

@@ -0,0 +1,46 @@
+inbound:
+  name: "b-inbound-server"
+  server:
+    type: "mux"
+    mux: "muxed"
+  client:
+    tcp:
+      serverAddress: "localhost:8233"
+outbound:
+  name: "b-outbound-server"
+  server:
+    tcp:
+      listenAddress: "0.0.0.0:6433"
+  client:
+    type: "mux"
+    mux: "muxed"
+mux:
+  - name: "muxed"
+    mode: "server"
+    server:
+      listenAddress: "0.0.0.0:6434"
+# shardCount:
+#   mode: "lcm"
+#   localShardCount: 3
+#   remoteShardCount: 2
+shardCount:
+  mode: "routing"
+  localShardCount: 3
+  remoteShardCount: 2
+profiling:
+  pprofAddress: "localhost:6071"
+memberlist:
+  enabled: true
+  nodeName: "proxy-node-b-2"
+  bindAddr: "127.0.0.1"
+  bindPort: 6435
+  joinAddrs:
+    - "localhost:6335"
+  proxyAddresses:
+    "proxy-node-b-1": "localhost:6333"
+    "proxy-node-b-2": "localhost:6433"
+  # # TCP-only configuration for restricted networks
+  # tcpOnly: true           # Use TCP transport only, disable UDP
+  # disableTCPPings: true   # Disable TCP pings for faster convergence
+  # probeTimeoutMs: 1000    # Longer timeout for network latency
+  # probeIntervalMs: 2000   # Less frequent probes to reduce network noise

develop/config/dynamic-config.yaml

@@ -20,10 +20,14 @@ history.ReplicationEnableUpdateWithNewTaskMerge:
 history.enableWorkflowExecutionTimeoutTimer:
   - value: true
 history.EnableReplicationTaskTieredProcessing:
-  - value: true
+  - value: false
 history.persistenceMaxQPS:
   - value: 100000
     constraints: {}
 frontend.persistenceMaxQPS:
   - value: 100000
-    constraints: {}
+    constraints: {}
+history.shardUpdateMinInterval:
+  - value: 1s
+history.ReplicationStreamSendEmptyTaskDuration:
+  - value: 10s

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/gogo/status v1.1.1
 	github.com/golang/mock v1.7.0-rc.1
 	github.com/grpc-ecosystem/go-grpc-middleware/providers/prometheus v1.1.0
+	github.com/hashicorp/memberlist v0.5.1
 	github.com/hashicorp/yamux v0.1.2
 	github.com/keilerkonzept/visit v1.1.1
 	github.com/pkg/errors v0.9.1
@@ -40,6 +41,7 @@ require (
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/exporter/metric v0.51.0 // indirect
 	github.com/GoogleCloudPlatform/opentelemetry-operations-go/internal/resourcemapping v0.51.0 // indirect
 	github.com/apache/thrift v0.21.0 // indirect
+	github.com/armon/go-metrics v0.0.0-20180917152333-f0300d1749da // indirect
 	github.com/aws/aws-sdk-go v1.55.6 // indirect
 	github.com/benbjohnson/clock v1.3.5 // indirect
 	github.com/beorn7/perks v1.0.1 // indirect
@@ -66,6 +68,7 @@ require (
 	github.com/golang-jwt/jwt/v4 v4.5.2 // indirect
 	github.com/golang/protobuf v1.5.4 // indirect
 	github.com/golang/snappy v1.0.0 // indirect
+	github.com/google/btree v0.0.0-20180813153112-4030bb1f1f0c // indirect
 	github.com/google/go-cmp v0.7.0 // indirect
 	github.com/google/s2a-go v0.1.9 // indirect
 	github.com/google/uuid v1.6.0 // indirect
@@ -76,6 +79,12 @@ require (
 	github.com/grpc-ecosystem/go-grpc-middleware/v2 v2.3.2 // indirect
 	github.com/grpc-ecosystem/grpc-gateway/v2 v2.26.3 // indirect
 	github.com/hailocab/go-hostpool v0.0.0-20160125115350-e80d13ce29ed // indirect
+	github.com/hashicorp/errwrap v1.0.0 // indirect
+	github.com/hashicorp/go-immutable-radix v1.0.0 // indirect
+	github.com/hashicorp/go-msgpack/v2 v2.1.1 // indirect
+	github.com/hashicorp/go-multierror v1.0.0 // indirect
+	github.com/hashicorp/go-sockaddr v1.0.0 // indirect
+	github.com/hashicorp/golang-lru v0.5.0 // indirect
 	github.com/hashicorp/golang-lru/v2 v2.0.7 // indirect
 	github.com/iancoleman/strcase v0.3.0 // indirect
 	github.com/jackc/pgpassfile v1.0.0 // indirect
@@ -89,6 +98,7 @@ require (
 	github.com/lib/pq v1.10.9 // indirect
 	github.com/mailru/easyjson v0.9.0 // indirect
 	github.com/mattn/go-isatty v0.0.20 // indirect
+	github.com/miekg/dns v1.1.26 // indirect
 	github.com/mitchellh/mapstructure v1.5.0 // indirect
 	github.com/munnerz/goautoneg v0.0.0-20191010083416-a7dc8b61c822 // indirect
 	github.com/ncruces/go-strftime v0.1.9 // indirect
@@ -106,6 +116,7 @@ require (
 	github.com/robfig/cron v1.2.0 // indirect
 	github.com/robfig/cron/v3 v3.0.1 // indirect
 	github.com/russross/blackfriday/v2 v2.1.0 // indirect
+	github.com/sean-/seed v0.0.0-20170313163322-e2103e2c3529 // indirect
 	github.com/sirupsen/logrus v1.9.3 // indirect
 	github.com/sony/gobreaker v1.0.0 // indirect
 	github.com/stretchr/objx v0.5.2 // indirect