Skip to content

feat(gateway): load balancing #370

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
bernie-g wants to merge 5 commits into
base: main
Choose a base branch
from
Open

feat(gateway): load balancing #370

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
90e25b4
feat(gateway): report active channel count for pool load balancing
bernie-g Aug 18, 2026
6210b36
fix(gateway): bound the load report request
bernie-g Aug 18, 2026
1d66d56
refactor(gateway): rename the load report endpoint to /metrics
bernie-g Aug 20, 2026
f5b073a
chore(deps): bump pion/stun to v3.1.5
bernie-g Aug 20, 2026
5e9ed35
chore(deps): sync the e2e module with the stun bump
bernie-g Aug 20, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions e2e/go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 4 additions & 4 deletions e2e/go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
4 changes: 2 additions & 2 deletions go.mod
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
8 changes: 4 additions & 4 deletions go.sum
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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=
Expand Down
21 changes: 21 additions & 0 deletions packages/api/api.go
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,7 @@
package api

import (
"context"
"encoding/base64"
"errors"
"fmt"
Expand Down Expand Up @@ -45,6 +46,7 @@ const (
operationCallExchangeRelayCertV1 = "CallExchangeRelayCertV1"
operationCallGatewayHeartBeatV1 = "CallGatewayHeartBeatV1"
operationCallGatewayHeartBeatV2 = "CallGatewayHeartBeatV2"
operationCallGatewayMetricsReportV2 = "CallGatewayMetricsReportV2"
operationCallBootstrapInstance = "CallBootstrapInstance"
operationCallRegisterInstanceRelay = "CallRegisterInstanceRelay"
operationCallRegisterOrgRelay = "CallRegisterOrgRelay"
Expand Down Expand Up @@ -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().
Expand Down
4 changes: 4 additions & 0 deletions packages/api/model.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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"`
Expand Down
55 changes: 55 additions & 0 deletions packages/gateway-v2/gateway.go
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -57,6 +57,16 @@ const (

const heartbeatInterval = 3 * time.Minute

// Reported over the same direct HTTP client as the heartbeat. The relay connection cannot carry this:
// the gateway only accepts channels on it, so it runs platform-to-gateway and has no reverse path.
// Kept off the heartbeat itself because that handler probes back through the relay and writes to the
// platform's database, which is far too costly at the cadence pool selection needs.
const metricsReportInterval = 10 * time.Second

// Must stay below the interval so a stalled endpoint cannot hold the reporting loop past the next
// tick. Without a bound the loop would also stop noticing cancellation and shutdown would hang.
const metricsReportTimeout = 5 * time.Second

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"
Expand Down Expand Up @@ -144,6 +154,11 @@ type Gateway struct {
mongoProxies map[string]*mongoProxyEntry
mongoProxiesMu sync.Mutex
pkcs11Module Pkcs11Module

// Every channel reaching this gateway is accepted in handleIncomingChannel, whoever opened it,
// so counting there is the only place that cannot be bypassed by a new caller. The platform uses
// it to pick the least loaded member of a gateway pool.
activeChannels atomic.Int64
}

// mongoProxyEntry holds a session-level MongoDB proxy with a ready signal.
Expand Down Expand Up @@ -384,6 +399,42 @@ func (g *Gateway) reapIdleSessions() {
}
}

// reportLoad publishes the gateway's active channel count so the platform can route new work to the

@x032205 x032205 Aug 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

The comment mentions "reportLoad" but this is the "sendMetricsReport" function

@x032205 x032205 Aug 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also this comment may be for "reportMetrics" instead of "sendMetricsReport"? Not 100% sure

@x032205 x032205 Aug 21, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Also maybe we should rename "reportMetrics" to "startMetricsReport" to be more clear

// least loaded member of a pool. Failures are not surfaced: a missed report only costs accuracy, and
// the platform falls back to its own view when a gateway stops reporting.
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})
}

func (g *Gateway) reportMetrics(ctx context.Context) {
go func() {
ticker := time.NewTicker(metricsReportInterval)
defer ticker.Stop()

var last int64 = -1
for {
select {
case <-ctx.Done():
return
case <-ticker.C:
count := g.activeChannels.Load()
// Still republish an unchanged count so the platform can tell a quiet gateway from
// one that has stopped reporting.
Comment thread
bernie-g marked this conversation as resolved.
if err := g.sendMetricsReport(ctx, count); err != nil {
log.Debug().Msgf("Load report failed: %v", err)

@x032205 x032205 Aug 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

An older self hosted platform has no metrics endpoint, so this retries every 10s forever. Maybe we should have some back-off mechanism that stops the report after X failures?

continue
}
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{}
Expand Down Expand Up @@ -575,6 +626,7 @@ func (g *Gateway) startHeartbeatOnce(ctx context.Context, errCh chan error) {
defer g.heartbeatMu.Unlock()
if !g.heartbeatStarted {
g.registerHeartBeat(ctx, errCh)
g.reportMetrics(ctx)
g.heartbeatStarted = true
}
}
Expand Down Expand Up @@ -884,6 +936,9 @@ func (g *Gateway) handleIncomingChannel(newChannel ssh.NewChannel) {
}
defer channel.Close()

g.activeChannels.Add(1)

@x032205 x032205 Aug 19, 2026

Copy link
Copy Markdown
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Claude:

This counter only goes down via the defer, so a leaked goroutine inflates it permanently and the gateway stops getting work. Could we reset it to 0 on each new relay connection?

defer g.activeChannels.Add(-1)

go ssh.DiscardRequests(requests)

// Create mTLS server configuration
Expand Down
Loading