feat(gateway): load balancing

e2e/go.mod

@@ -250,9 +250,9 @@ require (
 	github.com/pion/dtls/v3 v3.1.4 // indirect
 	github.com/pion/logging v0.2.4 // indirect
 	github.com/pion/randutil v0.1.0 // indirect
-	github.com/pion/stun/v3 v3.0.0 // indirect
+	github.com/pion/stun/v3 v3.1.5 // indirect
 	github.com/pion/transport/v3 v3.0.7 // indirect
-	github.com/pion/transport/v4 v4.0.1 // indirect
+	github.com/pion/transport/v4 v4.0.2 // indirect
 	github.com/pion/turn/v4 v4.0.0 // indirect
 	github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c // indirect
 	github.com/pkg/errors v0.9.1 // indirect

e2e/go.sum

@@ -857,12 +857,12 @@ github.com/pion/logging v0.2.4 h1:tTew+7cmQ+Mc1pTBLKH2puKsOvhm32dROumOZ655zB8=
 github.com/pion/logging v0.2.4/go.mod h1:DffhXTKYdNZU+KtJ5pyQDjvOAh/GsNSyv1lbkFbe3so=
 github.com/pion/randutil v0.1.0 h1:CFG1UdESneORglEsnimhUjf33Rwjubwj6xfiOXBa3mA=
 github.com/pion/randutil v0.1.0/go.mod h1:XcJrSMMbbMRhASFVOlj/5hQial/Y8oH/HVo7TBZq+j8=
-github.com/pion/stun/v3 v3.0.0 h1:4h1gwhWLWuZWOJIJR9s2ferRO+W3zA/b6ijOI6mKzUw=
-github.com/pion/stun/v3 v3.0.0/go.mod h1:HvCN8txt8mwi4FBvS3EmDghW6aQJ24T+y+1TKjB5jyU=
+github.com/pion/stun/v3 v3.1.5 h1:Y1FHlhaI6+4UoC5i/zQf4F7JvdZtB24/05oyy/GF1x8=
+github.com/pion/stun/v3 v3.1.5/go.mod h1:zRUghXSQU32Lx5orJsz3uYMkIihweXb3mu5gIns02fs=
 github.com/pion/transport/v3 v3.0.7 h1:iRbMH05BzSNwhILHoBoAPxoB9xQgOaJk+591KC9P1o0=
 github.com/pion/transport/v3 v3.0.7/go.mod h1:YleKiTZ4vqNxVwh77Z0zytYi7rXHl7j6uPLGhhz9rwo=
-github.com/pion/transport/v4 v4.0.1 h1:sdROELU6BZ63Ab7FrOLn13M6YdJLY20wldXW2Cu2k8o=
-github.com/pion/transport/v4 v4.0.1/go.mod h1:nEuEA4AD5lPdcIegQDpVLgNoDGreqM/YqmEx3ovP4jM=
+github.com/pion/transport/v4 v4.0.2 h1:ifYlPqNwsy6aKQ9y8yzxXlHae5431ZrH2avkD/Rn6Tk=
+github.com/pion/transport/v4 v4.0.2/go.mod h1:06hFI+jCFcok2X2MekVufNZ/uzNZXivGBPfviSVcjgM=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
 github.com/pkg/errors v0.8.0/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

go.mod

@@ -181,9 +181,9 @@ require (
 	github.com/pingcap/log v1.1.1-0.20241212030209-7e3ff8601a2a // indirect
 	github.com/pingcap/tidb/pkg/parser v0.0.0-20250421232622-526b2c79173d // indirect
 	github.com/pion/randutil v0.1.0 // indirect
-	github.com/pion/stun/v3 v3.0.0 // indirect
+	github.com/pion/stun/v3 v3.1.5 // indirect
 	github.com/pion/transport/v3 v3.0.7 // indirect
-	github.com/pion/transport/v4 v4.0.1 // indirect
+	github.com/pion/transport/v4 v4.0.2 // indirect
 	github.com/pmezard/go-difflib v1.0.1-0.20181226105442-5d4384ee4fb2 // indirect
 	github.com/rivo/uniseg v0.2.0 // indirect
 	github.com/shopspring/decimal v1.4.0 // indirect

go.sum

@@ -570,12 +570,12 @@ github.com/pion/logging v0.2.4 h1:tTew+7cmQ+Mc1pTBLKH2puKsOvhm32dROumOZ655zB8=
 github.com/pion/logging v0.2.4/go.mod h1:DffhXTKYdNZU+KtJ5pyQDjvOAh/GsNSyv1lbkFbe3so=
 github.com/pion/randutil v0.1.0 h1:CFG1UdESneORglEsnimhUjf33Rwjubwj6xfiOXBa3mA=
 github.com/pion/randutil v0.1.0/go.mod h1:XcJrSMMbbMRhASFVOlj/5hQial/Y8oH/HVo7TBZq+j8=
-github.com/pion/stun/v3 v3.0.0 h1:4h1gwhWLWuZWOJIJR9s2ferRO+W3zA/b6ijOI6mKzUw=
-github.com/pion/stun/v3 v3.0.0/go.mod h1:HvCN8txt8mwi4FBvS3EmDghW6aQJ24T+y+1TKjB5jyU=
+github.com/pion/stun/v3 v3.1.5 h1:Y1FHlhaI6+4UoC5i/zQf4F7JvdZtB24/05oyy/GF1x8=
+github.com/pion/stun/v3 v3.1.5/go.mod h1:zRUghXSQU32Lx5orJsz3uYMkIihweXb3mu5gIns02fs=
 github.com/pion/transport/v3 v3.0.7 h1:iRbMH05BzSNwhILHoBoAPxoB9xQgOaJk+591KC9P1o0=
 github.com/pion/transport/v3 v3.0.7/go.mod h1:YleKiTZ4vqNxVwh77Z0zytYi7rXHl7j6uPLGhhz9rwo=
-github.com/pion/transport/v4 v4.0.1 h1:sdROELU6BZ63Ab7FrOLn13M6YdJLY20wldXW2Cu2k8o=
-github.com/pion/transport/v4 v4.0.1/go.mod h1:nEuEA4AD5lPdcIegQDpVLgNoDGreqM/YqmEx3ovP4jM=
+github.com/pion/transport/v4 v4.0.2 h1:ifYlPqNwsy6aKQ9y8yzxXlHae5431ZrH2avkD/Rn6Tk=
+github.com/pion/transport/v4 v4.0.2/go.mod h1:06hFI+jCFcok2X2MekVufNZ/uzNZXivGBPfviSVcjgM=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c h1:+mdjkGKdHQG3305AYmdv1U2eRNDiU2ErMBj1gwrq8eQ=
 github.com/pkg/browser v0.0.0-20240102092130-5ac0b6a4141c/go.mod h1:7rwL4CYBLnjLxUqIJNnCWiEdr3bn6IUYi15bNlnbCCU=
 github.com/pkg/errors v0.8.1/go.mod h1:bwawxfHBFNV+L2hUp1rHADufV3IMtnDRdf1r5NINEl0=

packages/api/api.go

@@ -1,6 +1,7 @@
 package api
 
 import (
+	"context"
 	"encoding/base64"
 	"errors"
 	"fmt"
@@ -45,6 +46,7 @@ const (
 	operationCallExchangeRelayCertV1               = "CallExchangeRelayCertV1"
 	operationCallGatewayHeartBeatV1                = "CallGatewayHeartBeatV1"
 	operationCallGatewayHeartBeatV2                = "CallGatewayHeartBeatV2"
+	operationCallGatewayMetricsReportV2            = "CallGatewayMetricsReportV2"
 	operationCallBootstrapInstance                 = "CallBootstrapInstance"
 	operationCallRegisterInstanceRelay             = "CallRegisterInstanceRelay"
 	operationCallRegisterOrgRelay                  = "CallRegisterOrgRelay"
@@ -847,6 +849,25 @@ func CallGatewayHeartBeatV2(httpClient *resty.Client, request GatewayHeartbeatRe
 	return nil
 }
 
+func CallGatewayMetricsReportV2(ctx context.Context, httpClient *resty.Client, request GatewayMetricsReportRequest) error {
+	response, err := httpClient.
+		R().
+		SetContext(ctx).
+		SetHeader("User-Agent", USER_AGENT).
+		SetBody(request).
+		Post(fmt.Sprintf("%v/v2/gateways/metrics", config.INFISICAL_URL))
+
+	if err != nil {
+		return NewGenericRequestError(operationCallGatewayMetricsReportV2, err)
+	}
+
+	if response.IsError() {
+		return NewAPIErrorWithResponse(operationCallGatewayMetricsReportV2, response, nil)
+	}
+
+	return nil
+}
+
 func CallOrgRelayHeartBeat(httpClient *resty.Client, request RelayHeartbeatRequest) error {
 	response, err := httpClient.
 		R().

packages/api/model.go

@@ -1084,6 +1084,10 @@ type GatewayHeartbeatRequest struct {
 	Capabilities map[string]any `json:"capabilities,omitempty"`
 }
 
+type GatewayMetricsReportRequest struct {
+	ActiveChannels int64 `json:"activeChannels"`
+}
+
 type RelayLoginRequest struct {
 	Method            string `json:"method"`
 	Token             string `json:"token,omitempty"`

packages/gateway-v2/gateway.go

@@ -57,6 +57,16 @@ const (
 
 const heartbeatInterval = 3 * time.Minute
 
+// Kept off the heartbeat because that handler probes back through the relay and writes to the
+// platform's database, which is far too costly at this cadence.
+const metricsReportInterval = 10 * time.Second
+
+// Below the interval, so a stalled endpoint cannot hold the loop past the next tick or block shutdown.
+const metricsReportTimeout = 5 * time.Second
+
+const metricsReportFailuresBeforeBackoff = 5
+const metricsReportBackoff = 5 * time.Minute
+
 const GATEWAY_ROUTING_INFO_OID = "1.3.6.1.4.1.12345.100.1"
 const GATEWAY_ACTOR_OID = "1.3.6.1.4.1.12345.100.2"
 const PAM_INFO_OID = "1.3.6.1.4.1.12345.100.3"
@@ -144,6 +154,11 @@ type Gateway struct {
 	mongoProxies   map[string]*mongoProxyEntry
 	mongoProxiesMu sync.Mutex
 	pkcs11Module   Pkcs11Module
+
+	// Counted in the relay's channel-receive loop, which no caller can bypass.
+	activeChannels atomic.Int64
+	// Bumped per relay connection, so a handler cannot release a count it did not acquire.
+	channelGeneration atomic.Int64
 }
 
 // mongoProxyEntry holds a session-level MongoDB proxy with a ready signal.
@@ -384,6 +399,53 @@ func (g *Gateway) reapIdleSessions() {
 	}
 }
 
+func (g *Gateway) sendMetricsReport(ctx context.Context, count int64) error {
+	reqCtx, cancel := context.WithTimeout(ctx, metricsReportTimeout)
+	defer cancel()
+	return api.CallGatewayMetricsReportV2(reqCtx, g.httpClient, api.GatewayMetricsReportRequest{ActiveChannels: count})
+}
+
+// A gateway that stops reporting takes its whole pool off load-aware selection.
+func (g *Gateway) startMetricsReport(ctx context.Context) {
+	go func() {
+		// Report immediately: until one lands the pool has nothing to compare and selects at random.
+		delay := time.Duration(0)
+		failures := 0
+
+		var last int64 = -1
+		for {
+			select {
+			case <-ctx.Done():
+				return
+			case <-time.After(delay):
+			}
+			delay = metricsReportInterval
+
+			count := g.activeChannels.Load()
+			// Republish unchanged, so a quiet gateway is distinguishable from a silent one.
+			if err := g.sendMetricsReport(ctx, count); err != nil {
+				failures++
+				if failures == metricsReportFailuresBeforeBackoff {
+					log.Warn().Err(err).Msgf("Metrics report failing; backing off to %s. Pools containing this gateway will select at random until it succeeds", metricsReportBackoff)
+				}
+				if failures >= metricsReportFailuresBeforeBackoff {
+					delay = metricsReportBackoff
+				}
+				continue
+			}
+
+			if failures >= metricsReportFailuresBeforeBackoff {
+				log.Info().Msg("Metrics report recovered")
+			}
+			failures = 0
+			if count != last {
+				log.Debug().Msgf("Reported %d active channels", count)
+				last = count
+			}
+		}
+	}()
+}
+
 func (g *Gateway) registerHeartBeat(ctx context.Context, errCh chan error) {
 	sendHeartbeat := func() error {
 		capabilities := map[string]any{}
@@ -575,6 +637,7 @@ func (g *Gateway) startHeartbeatOnce(ctx context.Context, errCh chan error) {
 	defer g.heartbeatMu.Unlock()
 	if !g.heartbeatStarted {
 		g.registerHeartBeat(ctx, errCh)
+		g.startMetricsReport(ctx)
 		g.heartbeatStarted = true
 	}
 }
@@ -650,6 +713,10 @@ func (g *Gateway) handleConnection(client *ssh.Client) error {
 		client.Close()
 	}()
 
+	// Channels do not outlive their connection, so anything still counted is a handler that hung.
+	generation := g.channelGeneration.Add(1)
+	g.activeChannels.Store(0)
+
 	// Handle incoming channels from the server
 	channels := client.HandleChannelOpen("direct-tcpip")
 	if channels == nil {
@@ -693,7 +760,9 @@ func (g *Gateway) handleConnection(client *ssh.Client) error {
 				log.Info().Msg("SSH channels closed")
 				return nil
 			}
-			go g.handleIncomingChannel(newChannel)
+			// Counted here, not in the handler: its goroutine could load a later generation.
+			g.activeChannels.Add(1)
+			go g.handleIncomingChannel(newChannel, generation)
 		}
 	}
 }
@@ -876,7 +945,25 @@ func (g *Gateway) validateHostCertificate(cert *ssh.Certificate, hostname string
 	return nil
 }
 
-func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) {
+func (g *Gateway) releaseChannel(generation int64) {
+	// A handler outliving its connection would otherwise decrement a count it never contributed to.
+	if g.channelGeneration.Load() != generation {
+		return
+	}
+	for {
+		current := g.activeChannels.Load()
+		if current <= 0 {
+			return
+		}
+		if g.activeChannels.CompareAndSwap(current, current-1) {
+			return
+		}
+	}
+}
+
+func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel, generation int64) {
+	defer g.releaseChannel(generation)
+
 	channel, requests, err := newChannel.Accept()
 	if err != nil {
 		log.Info().Msgf("Failed to accept channel: %v", err)

packages/gateway-v2/release_channel_test.go

@@ -0,0 +1,54 @@
+package gatewayv2
+
+import (
+	"sync"
+	"testing"
+)
+
+func TestReleaseChannelDecrements(t *testing.T) {
+	g := &Gateway{}
+	g.activeChannels.Add(3)
+	g.releaseChannel(g.channelGeneration.Load())
+	if got := g.activeChannels.Load(); got != 2 {
+		t.Fatalf("count is %d, want 2", got)
+	}
+}
+
+func TestReleaseChannelIgnoresPreviousGeneration(t *testing.T) {
+	g := &Gateway{}
+	stale := g.channelGeneration.Load()
+	g.activeChannels.Add(2)
+
+	// A relay reconnect: generation bumps and the count resets.
+	g.channelGeneration.Add(1)
+	g.activeChannels.Store(0)
+	g.activeChannels.Add(1)
+
+	g.releaseChannel(stale)
+	g.releaseChannel(stale)
+
+	if got := g.activeChannels.Load(); got != 1 {
+		t.Fatalf("stale releases changed the count to %d, want 1", got)
+	}
+}
+
+func TestReleaseChannelNeverGoesNegative(t *testing.T) {
+	g := &Gateway{}
+	generation := g.channelGeneration.Load()
+	g.activeChannels.Add(2)
+	g.activeChannels.Store(0)
+
+	var wg sync.WaitGroup
+	for i := 0; i < 2; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			g.releaseChannel(generation)
+		}()
+	}
+	wg.Wait()
+
+	if got := g.activeChannels.Load(); got != 0 {
+		t.Fatalf("count went to %d, want 0", got)
+	}
+}