diff --git a/e2e/go.mod b/e2e/go.mod index ed358a56..0fc80dc2 100644 --- a/e2e/go.mod +++ b/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 diff --git a/e2e/go.sum b/e2e/go.sum index e873f8fb..9ed8bee5 100644 --- a/e2e/go.sum +++ b/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= diff --git a/go.mod b/go.mod index e5acef0f..4dcfc337 100644 --- a/go.mod +++ b/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 diff --git a/go.sum b/go.sum index 993a0108..c69d3a3a 100644 --- a/go.sum +++ b/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= diff --git a/packages/api/api.go b/packages/api/api.go index 9bce7e93..1707cfd2 100644 --- a/packages/api/api.go +++ b/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(). diff --git a/packages/api/model.go b/packages/api/model.go index ec9aeddb..3188f73c 100644 --- a/packages/api/model.go +++ b/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"` diff --git a/packages/gateway-v2/gateway.go b/packages/gateway-v2/gateway.go index 9fdae907..52974335 100644 --- a/packages/gateway-v2/gateway.go +++ b/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,9 @@ type Gateway struct { mongoProxies map[string]*mongoProxyEntry mongoProxiesMu sync.Mutex pkcs11Module Pkcs11Module + + // Counted in handleIncomingChannel, the one place no caller can bypass. + activeChannels atomic.Int64 } // mongoProxyEntry holds a session-level MongoDB proxy with a ready signal. @@ -384,6 +397,52 @@ 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() { + delay := metricsReportInterval + failures := 0 + + var last int64 = -1 + for { + select { + case <-ctx.Done(): + return + case <-time.After(delay): + } + + 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 + delay = metricsReportInterval + 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 +634,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 +710,9 @@ 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. + g.activeChannels.Store(0) + // Handle incoming channels from the server channels := client.HandleChannelOpen("direct-tcpip") if channels == nil { @@ -876,6 +939,19 @@ func (g *Gateway) validateHostCertificate(cert *ssh.Certificate, hostname string return nil } +// Floored: a handler from a previous connection must not decrement past the reset. +func (g *Gateway) releaseChannel() { + for { + current := g.activeChannels.Load() + if current <= 0 { + return + } + if g.activeChannels.CompareAndSwap(current, current-1) { + return + } + } +} + func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) { channel, requests, err := newChannel.Accept() if err != nil { @@ -884,6 +960,9 @@ func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) { } defer channel.Close() + g.activeChannels.Add(1) + defer g.releaseChannel() + go ssh.DiscardRequests(requests) // Create mTLS server configuration diff --git a/packages/gateway-v2/release_channel_test.go b/packages/gateway-v2/release_channel_test.go new file mode 100644 index 00000000..2c35a41f --- /dev/null +++ b/packages/gateway-v2/release_channel_test.go @@ -0,0 +1,38 @@ +package gatewayv2 + +import ( + "sync" + "testing" +) + +func TestReleaseChannelNeverGoesNegative(t *testing.T) { + g := &Gateway{} + g.activeChannels.Add(1) + g.activeChannels.Add(1) + + // Reset with handlers still in flight, as on a relay reconnect. + g.activeChannels.Store(0) + + var wg sync.WaitGroup + for i := 0; i < 2; i++ { + wg.Add(1) + go func() { + defer wg.Done() + g.releaseChannel() + }() + } + wg.Wait() + + if got := g.activeChannels.Load(); got != 0 { + t.Fatalf("count went to %d, want 0", got) + } +} + +func TestReleaseChannelDecrements(t *testing.T) { + g := &Gateway{} + g.activeChannels.Add(3) + g.releaseChannel() + if got := g.activeChannels.Load(); got != 2 { + t.Fatalf("count is %d, want 2", got) + } +}