reversetunnel: support tracker discovery ttl and disconnects (#65419)

* reversetunnel: support tracker discovery ttl and disconnects

This adds agent side support for tracker ttl to be configured
per proxy over discovery requests.

This gives us flexibility of lowering the tracker ttl server side.

Lowering the ttl has the benefit of dropping proxies from the tracker
sooner allowing the agent to attempt to reconnect to proxies sooner.

This change also adds agent disconnects. This only occurs if a disconnect
threshold is configured in the cluster network config.

A proxy is only ever considered for disconnect if its in the desired
set and the agent has more connections than configured.

The agent will disconnect from proxies with the highest smoothed rtt.

The agent will not disconnect from proxies during a rollout where
the group/generation is changed or from unhealthy proxies which
are no longer tracked.

* fix lint

* stop agent in background and use testsync

* remove gogoproto and move to disconnect_threshold_seconds

* remove fake clock from tracker

* simplify conditions and nesting

* update comment about proxy disconnects

* generate terraform

* move now inside lock

* add else for smoothedRTT calculation

* rename disconnect to disconnectCandidate for clarity

* add comment for tryDisconnect usage

* rearange some checks in tryDisconnect

* fix test comment

* preserve cloud managed cnf fields on update/upsert
This commit is contained in:
david boslee
2026-06-12 18:22:01 +00:00
committed by GitHub
parent 36166db5f1
commit a44ab9b510
16 changed files with 3054 additions and 2335 deletions
@@ -2623,6 +2623,9 @@ message AgentMeshTunnelStrategy {}
// ProxyPeeringTunnelStrategy requires reverse tunnels to dial a fixed number of proxies.
message ProxyPeeringTunnelStrategy {
int64 AgentConnectionCount = 1 [(gogoproto.jsontag) = "agent_connection_count,omitempty"];
// disconnect_threshold_seconds is the duration without connectivity or topology changes
// before an agent pool may disconnect excess connections.
uint32 disconnect_threshold_seconds = 2;
}
// ProxyListenerMode represents the cluster proxy listener mode.
+10
View File
@@ -52,6 +52,16 @@ func TestTunnelStrategyMarshalling(t *testing.T) {
},
}},
},
{
json: `{"agent_connection_count":2,"disconnect_threshold_seconds":60,"type":"proxy_peering"}`,
yaml: "agent_connection_count: 2\ndisconnect_threshold_seconds: 60\ntype: proxy_peering\n",
strategy: &TunnelStrategyV1{Strategy: &TunnelStrategyV1_ProxyPeering{
ProxyPeering: &ProxyPeeringTunnelStrategy{
AgentConnectionCount: 2,
DisconnectThresholdSeconds: 60,
},
}},
},
}
for _, tc := range tests {
+2318 -2286
View File
File diff suppressed because it is too large Load Diff
@@ -74,4 +74,5 @@ Optional:
Optional:
- `agent_connection_count` (Number)
- `disconnect_threshold_seconds` (Number) disconnect_threshold_seconds is the duration without connectivity or topology changes before an agent pool may disconnect excess connections.
@@ -93,3 +93,4 @@ Optional:
Optional:
- `agent_connection_count` (Number)
- `disconnect_threshold_seconds` (Number) disconnect_threshold_seconds is the duration without connectivity or topology changes before an agent pool may disconnect excess connections.
@@ -2237,11 +2237,18 @@ func GenSchemaClusterNetworkingConfigV2(ctx context.Context) (github_com_hashico
Optional: true,
},
"proxy_peering": {
Attributes: github_com_hashicorp_terraform_plugin_framework_tfsdk.SingleNestedAttributes(map[string]github_com_hashicorp_terraform_plugin_framework_tfsdk.Attribute{"agent_connection_count": {
Description: "",
Optional: true,
Type: github_com_hashicorp_terraform_plugin_framework_types.Int64Type,
}}),
Attributes: github_com_hashicorp_terraform_plugin_framework_tfsdk.SingleNestedAttributes(map[string]github_com_hashicorp_terraform_plugin_framework_tfsdk.Attribute{
"agent_connection_count": {
Description: "",
Optional: true,
Type: github_com_hashicorp_terraform_plugin_framework_types.Int64Type,
},
"disconnect_threshold_seconds": {
Description: "disconnect_threshold_seconds is the duration without connectivity or topology changes before an agent pool may disconnect excess connections.",
Optional: true,
Type: github_com_hashicorp_terraform_plugin_framework_types.Int64Type,
},
}),
Description: "",
Optional: true,
},
@@ -24646,6 +24653,23 @@ func CopyClusterNetworkingConfigV2FromTerraform(_ context.Context, tf github_com
}
}
}
{
a, ok := tf.Attrs["disconnect_threshold_seconds"]
if !ok {
diags.Append(attrReadMissingDiag{"ClusterNetworkingConfigV2.Spec.TunnelStrategy.ProxyPeering.disconnect_threshold_seconds"})
} else {
v, ok := a.(github_com_hashicorp_terraform_plugin_framework_types.Int64)
if !ok {
diags.Append(attrReadConversionFailureDiag{"ClusterNetworkingConfigV2.Spec.TunnelStrategy.ProxyPeering.disconnect_threshold_seconds", "github.com/hashicorp/terraform-plugin-framework/types.Int64"})
} else {
var t uint32
if !v.Null && !v.Unknown {
t = uint32(v.Value)
}
obj.DisconnectThresholdSeconds = t
}
}
}
}
}
}
@@ -25310,6 +25334,28 @@ func CopyClusterNetworkingConfigV2ToTerraform(ctx context.Context, obj *github_c
tf.Attrs["agent_connection_count"] = v
}
}
{
t, ok := tf.AttrTypes["disconnect_threshold_seconds"]
if !ok {
diags.Append(attrWriteMissingDiag{"ClusterNetworkingConfigV2.Spec.TunnelStrategy.ProxyPeering.disconnect_threshold_seconds"})
} else {
v, ok := tf.Attrs["disconnect_threshold_seconds"].(github_com_hashicorp_terraform_plugin_framework_types.Int64)
if !ok {
i, err := t.ValueFromTerraform(ctx, github_com_hashicorp_terraform_plugin_go_tftypes.NewValue(t.TerraformType(ctx), nil))
if err != nil {
diags.Append(attrWriteGeneralError{"ClusterNetworkingConfigV2.Spec.TunnelStrategy.ProxyPeering.disconnect_threshold_seconds", err})
}
v, ok = i.(github_com_hashicorp_terraform_plugin_framework_types.Int64)
if !ok {
diags.Append(attrWriteConversionFailureDiag{"ClusterNetworkingConfigV2.Spec.TunnelStrategy.ProxyPeering.disconnect_threshold_seconds", "github.com/hashicorp/terraform-plugin-framework/types.Int64"})
}
v.Null = int64(obj.DisconnectThresholdSeconds) == 0
}
v.Value = int64(obj.DisconnectThresholdSeconds)
v.Unknown = false
tf.Attrs["disconnect_threshold_seconds"] = v
}
}
}
v.Unknown = false
tf.Attrs["proxy_peering"] = v
@@ -501,6 +501,7 @@ func (s *Service) UpdateClusterNetworkingConfig(ctx context.Context, req *cluste
}
newCfg := req.GetClusterNetworkConfig()
PreserveCloudClusterNetworkConfig(*authzCtx, s.modules, newCfg, oldCfg)
if err := ValidateCloudNetworkConfigUpdate(*authzCtx, s.modules, newCfg, oldCfg); err != nil {
return nil, trace.Wrap(err)
@@ -566,6 +567,7 @@ func (s *Service) UpsertClusterNetworkingConfig(ctx context.Context, req *cluste
}
newCfg := req.GetClusterNetworkConfig()
PreserveCloudClusterNetworkConfig(*authzCtx, s.modules, newCfg, oldCfg)
if err := ValidateCloudNetworkConfigUpdate(*authzCtx, s.modules, newCfg, oldCfg); err != nil {
return nil, trace.Wrap(err)
@@ -698,8 +700,13 @@ func ValidateCloudNetworkConfigUpdate(authzCtx authz.Context, modules modules.Mo
oldts := oldConfig.GetProxyPeeringTunnelStrategy()
newts := newConfig.GetProxyPeeringTunnelStrategy()
if oldts != nil && newts != nil && oldts.AgentConnectionCount != newts.AgentConnectionCount {
return trace.BadParameter(cloudUpdateFailureMsg, "agent_connection_count")
if oldts != nil && newts != nil {
if oldts.AgentConnectionCount != newts.AgentConnectionCount {
return trace.BadParameter(cloudUpdateFailureMsg, "agent_connection_count")
}
if oldts.DisconnectThresholdSeconds != newts.DisconnectThresholdSeconds {
return trace.BadParameter(cloudUpdateFailureMsg, "disconnect_threshold")
}
}
if newConfig.GetKeepAliveInterval() != oldConfig.GetKeepAliveInterval() {
@@ -713,6 +720,27 @@ func ValidateCloudNetworkConfigUpdate(authzCtx authz.Context, modules modules.Mo
return nil
}
// PreserveCloudClusterNetworkConfig preserves cloud-managed fields when an older
// supported client attempts to update the cluster network config without knowing
// a field exists. This should only be used to preserve new fields that must only
// be configured by cloud.
func PreserveCloudClusterNetworkConfig(authzCtx authz.Context, modules modules.Modules, newConfig, oldConfig types.ClusterNetworkingConfig) {
if authz.HasBuiltinRole(authzCtx, string(types.RoleAdmin)) {
return
}
if !modules.Features().Cloud {
return
}
oldts := oldConfig.GetProxyPeeringTunnelStrategy()
newts := newConfig.GetProxyPeeringTunnelStrategy()
if oldts != nil && newts != nil {
if oldts.DisconnectThresholdSeconds != 0 && newts.DisconnectThresholdSeconds == 0 {
newts.DisconnectThresholdSeconds = oldts.DisconnectThresholdSeconds
}
}
}
// GetSessionRecordingConfig returns the locally cached networking configuration.
func (s *Service) GetSessionRecordingConfig(ctx context.Context, _ *clusterconfigpb.GetSessionRecordingConfigRequest) (*types.SessionRecordingConfigV2, error) {
authzCtx, err := s.authorizer.Authorize(ctx)
@@ -736,6 +736,87 @@ func TestUpdateClusterNetworkingConfig(t *testing.T) {
require.True(t, trace.IsBadParameter(err), "got (%v), expected cloud feature to prevent updating agent connection count", err)
},
},
{
name: "cloud disconnect threshold count",
authorizer: authz.AuthorizerFunc(func(ctx context.Context) (*authz.Context, error) {
return &authz.Context{
Checker: fakeChecker{
rules: map[string][]string{types.KindClusterNetworkingConfig: {types.VerbUpdate}},
},
AdminActionAuthState: authz.AdminActionAuthMFAVerified,
Identity: authz.RemoteUser{},
}, nil
}),
cnc: func() types.ClusterNetworkingConfig {
cnc := types.DefaultClusterNetworkingConfig()
cnc.SetTunnelStrategy(&types.TunnelStrategyV1{
Strategy: &types.TunnelStrategyV1_ProxyPeering{
ProxyPeering: types.DefaultProxyPeeringTunnelStrategy(),
},
})
return cnc
}(),
config: func(p types.ClusterNetworkingConfig) {
ts := p.GetProxyPeeringTunnelStrategy()
ts.DisconnectThresholdSeconds = 60
p.SetTunnelStrategy(&types.TunnelStrategyV1{
Strategy: &types.TunnelStrategyV1_ProxyPeering{
ProxyPeering: ts,
},
})
},
modules: &modulestest.Modules{
TestBuildType: modules.BuildEnterprise,
TestFeatures: modules.Features{
Cloud: true,
},
},
assertion: func(t *testing.T, updated types.ClusterNetworkingConfig, err error) {
require.True(t, trace.IsBadParameter(err), "got (%v), expected cloud feature to prevent updating disconnect threshold", err)
},
},
{
name: "cloud preserve disconnect threshold",
authorizer: authz.AuthorizerFunc(func(ctx context.Context) (*authz.Context, error) {
return &authz.Context{
Checker: fakeChecker{
rules: map[string][]string{types.KindClusterNetworkingConfig: {types.VerbUpdate}},
},
AdminActionAuthState: authz.AdminActionAuthMFAVerified,
Identity: authz.RemoteUser{},
}, nil
}),
cnc: func() types.ClusterNetworkingConfig {
cnc := types.DefaultClusterNetworkingConfig()
peering := types.DefaultProxyPeeringTunnelStrategy()
peering.DisconnectThresholdSeconds = 30
cnc.SetTunnelStrategy(&types.TunnelStrategyV1{
Strategy: &types.TunnelStrategyV1_ProxyPeering{
ProxyPeering: peering,
},
})
return cnc
}(),
config: func(p types.ClusterNetworkingConfig) {
ts := p.GetProxyPeeringTunnelStrategy()
ts.DisconnectThresholdSeconds = 0
p.SetTunnelStrategy(&types.TunnelStrategyV1{
Strategy: &types.TunnelStrategyV1_ProxyPeering{
ProxyPeering: ts,
},
})
},
modules: &modulestest.Modules{
TestBuildType: modules.BuildEnterprise,
TestFeatures: modules.Features{
Cloud: true,
},
},
assertion: func(t *testing.T, updated types.ClusterNetworkingConfig, err error) {
require.NoError(t, err)
require.Equal(t, uint32(30), updated.GetProxyPeeringTunnelStrategy().DisconnectThresholdSeconds)
},
},
}
for _, test := range cases {
+27 -4
View File
@@ -210,6 +210,9 @@ type agent struct {
drainWG sync.WaitGroup
// proxySigner is used to sign PROXY headers for securely propagating client IP address
proxySigner multiplexer.PROXYHeaderSigner
// smoothedRTT estimates the roundtrip time for keep alive requests using
// rolling average weighted towards more recent measurements.
smoothedRTT time.Duration
}
// newAgent intializes a reverse tunnel agent.
@@ -249,6 +252,15 @@ func (a *agent) GetProxyID() (string, bool) {
return proxyIDFromPrincipals(a.client.Principals())
}
func (a *agent) RTT() (time.Duration, bool) {
a.mu.Lock()
defer a.mu.Unlock()
if a.smoothedRTT == 0 {
return 0, false
}
return a.smoothedRTT, true
}
// proxyIDFromPrincipals gets the proxy id from a list of principals.
func proxyIDFromPrincipals(principals []string) (string, bool) {
if len(principals) == 0 {
@@ -538,7 +550,7 @@ func (a *agent) handleGlobalRequests(ctx context.Context, requests <-chan *ssh.R
// handleDrainChannels handles channels that should be stopped when the agent is draining.
func (a *agent) handleDrainChannels(drainWGDone func()) error {
ticker := time.NewTicker(a.keepAlive)
ticker := a.clock.NewTicker(a.keepAlive)
defer ticker.Stop()
drainCtxDone := a.drainCtx.Done()
@@ -557,7 +569,7 @@ func (a *agent) handleDrainChannels(drainWGDone func()) error {
// for good measure
ticker.Stop()
// Send ping over heartbeat channel.
case <-ticker.C:
case <-ticker.Chan():
if a.drainCtx.Err() != nil {
continue
}
@@ -636,17 +648,18 @@ func (a *agent) handleChannels() error {
// be enabled.
func (a *agent) sendKeepalives() error {
first := true
ticker := time.NewTimer(0)
ticker := a.clock.NewTimer(0)
defer ticker.Stop()
for {
select {
case <-a.ctx.Done():
return nil
case <-ticker.C:
case <-ticker.Chan():
}
const wantReplyTrue = true
now := a.clock.Now()
_, _, err := a.client.SendRequest(a.ctx, teleport.KeepAliveReqType, wantReplyTrue, nil)
ticker.Reset(retryutils.SeventhJitter(a.keepAlive))
if err != nil {
@@ -663,6 +676,16 @@ func (a *agent) sendKeepalives() error {
return trace.Wrap(err, "failed to send keepalive request")
}
a.mu.Lock()
rtt := a.clock.Since(now)
if a.smoothedRTT == 0 {
a.smoothedRTT = rtt
} else {
a.smoothedRTT = (7*a.smoothedRTT + rtt) / 8
}
a.logger.DebugContext(a.ctx, "Computed new SRTT", "srtt", a.smoothedRTT.String())
a.mu.Unlock()
if !first {
continue
}
+11
View File
@@ -84,3 +84,14 @@ func (s *agentStore) proxyIDs() []string {
}
return ids
}
func (s *agentStore) getByProxyID(proxyID string) (Agent, bool) {
s.mu.Lock()
defer s.mu.Unlock()
for _, agent := range s.agents {
if id, ok := agent.GetProxyID(); ok && id == proxyID {
return agent, true
}
}
return nil, false
}
+32 -1
View File
@@ -25,9 +25,18 @@ import (
"github.com/stretchr/testify/require"
)
type testStoreAgent struct {
Agent
proxyID string
}
func (a *testStoreAgent) GetProxyID() (string, bool) {
return a.proxyID, true
}
func TestAgentStoreRace(t *testing.T) {
store := newAgentStore()
agents := []*agent{{}, {}, {}, {}, {}}
agents := []*testStoreAgent{{}, {}, {}, {}, {}}
wg := &sync.WaitGroup{}
for i := range agents {
@@ -52,3 +61,25 @@ func TestAgentStoreRace(t *testing.T) {
wg.Wait()
}
func TestAgentStoreGetByProxyID(t *testing.T) {
store := newAgentStore()
first := &testStoreAgent{proxyID: "proxy-1"}
second := &testStoreAgent{proxyID: "proxy-2"}
store.add(first)
store.add(second)
got, ok := store.getByProxyID("proxy-2")
require.True(t, ok)
require.Same(t, second, got)
got, ok = store.getByProxyID("proxy-1")
require.True(t, ok)
require.Same(t, first, got)
got, ok = store.getByProxyID("missing")
require.False(t, ok)
require.Nil(t, got)
}
+42
View File
@@ -29,6 +29,7 @@ import (
"github.com/google/uuid"
"github.com/gravitational/trace"
"github.com/jonboulle/clockwork"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"golang.org/x/crypto/ssh"
@@ -378,6 +379,47 @@ func TestAgentStart(t *testing.T) {
require.Equal(t, AgentClosed, agent.GetState())
}
func TestAgentSmoothedRTT(t *testing.T) {
clock := clockwork.NewFakeClock()
agent, client := testAgent(t, agentConfig{
staleConnTimeoutDisabled: true,
keepAlive: time.Minute,
clock: clock,
})
ctx, cancel := context.WithCancel(t.Context())
t.Cleanup(cancel)
agent.ctx = ctx
agent.cancel = cancel
agent.client = client
count := &atomic.Int32{}
client.MockSendRequest = func(ctx context.Context, name string, wantReply bool, payload []byte) (bool, []byte, error) {
switch count.Add(1) {
case 1:
clock.Advance(20 * time.Millisecond)
case 2:
clock.Advance(100 * time.Millisecond)
default:
return false, nil, trace.Errorf("err")
}
return true, nil, nil
}
go agent.sendKeepalives()
require.NoError(t, clock.BlockUntilContext(ctx, 1))
clock.Advance(time.Minute)
require.NoError(t, clock.BlockUntilContext(ctx, 1))
rtt, ok := agent.RTT()
require.True(t, ok)
require.Equal(t, 2, int(count.Load()))
assert.Equal(t, 30*time.Millisecond, rtt)
}
func TestAgentStateTransitions(t *testing.T) {
tests := []struct {
start AgentState
+110 -5
View File
@@ -79,6 +79,9 @@ type AgentPool struct {
active *agentStore
tracker *track.Tracker
// lastConnectivityChange is updated on any agent event.
lastConnectivityChange time.Time
// runtimeConfig contains dynamic configuration values.
runtimeConfig *agentPoolRuntimeConfig
@@ -184,6 +187,9 @@ type Agent interface {
GetState() AgentState
// GetProxyID returns the proxy id of the proxy the agent is connected to.
GetProxyID() (string, bool)
// RTT returns a estimated roundtrip time. If the bool is false then an
// estimate has not yet been computed.
RTT() (time.Duration, bool)
}
// NewAgentPool returns new instance of the agent pool.
@@ -211,7 +217,8 @@ func NewAgentPool(ctx context.Context, config AgentPoolConfig) (*AgentPool, erro
"target_cluster", config.Cluster,
"local_cluster", config.LocalCluster,
),
runtimeConfig: newAgentPoolRuntimeConfig(config.InsecureMode),
runtimeConfig: newAgentPoolRuntimeConfig(config.InsecureMode),
lastConnectivityChange: config.Clock.Now(),
}
pool.runtimeConfig.isRemoteCluster = pool.IsRemoteCluster
@@ -291,6 +298,7 @@ func (p *AgentPool) run() error {
} else {
p.wg.Add(1)
p.active.add(agent)
p.lastConnectivityChange = p.Clock.Now()
p.updateConnectedProxies()
}
@@ -369,7 +377,7 @@ func (p *AgentPool) processEvents(ctx context.Context, events <-chan Agent) erro
case <-ctx.Done():
return trace.Wrap(ctx.Err())
case agent := <-events:
p.handleEvent(ctx, agent)
p.handleSyncEvent(ctx, agent)
continue
default:
}
@@ -397,14 +405,93 @@ func (p *AgentPool) waitForLease(ctx context.Context, events <-chan Agent) (*tra
select {
case <-ctx.Done():
case <-t.C:
p.tryDisconnect(ctx)
case agent := <-events:
p.handleEvent(ctx, agent)
p.handleSyncEvent(ctx, agent)
}
}
return nil, trace.Wrap(ctx.Err())
}
// tryDisconnect disconnects an agent if all the following are true:
// 1. A disconnect threshold has been configured in the clusternetworkconfig.
// 2. There have been no connectivity changes within the configured threshold
// 3. There have been no tracker topology changes within the configured threshold
// 4. There are more active agents than the configured connection count.
// This should be used synchonously within the main agentpool loop.
func (p *AgentPool) tryDisconnect(ctx context.Context) {
disconnectThreshold := p.runtimeConfig.getDisconnectThreshold()
if disconnectThreshold <= 0 {
return
}
proxyIDs := p.active.proxyIDs()
snapshot := p.tracker.Snapshot()
if snapshot.ConnectionCount == 0 {
return
}
now := p.Clock.Now()
if !p.lastConnectivityChange.Add(disconnectThreshold).Before(now) {
return
}
if !snapshot.LastTopologyChange.Add(disconnectThreshold).Before(now) {
return
}
var activeDesiredProxies int
for _, proxyID := range proxyIDs {
if desired := snapshot.Proxies[proxyID]; desired {
activeDesiredProxies++
}
}
if activeDesiredProxies <= snapshot.ConnectionCount {
return
}
var (
maxRTT time.Duration
disconnectCandidate Agent
)
for _, proxyID := range proxyIDs {
// Only active desired proxies are considered for disconnect.
// We keep undesired connections to provide availability during termination
// or recovery after failover. These connections belong to either:
// - unhealthy proxies which we expect to terminate or recover and
// re-enter the desired set.
// - older generation proxies we expect to terminate as part of a proxy
// rollout.
if desired := snapshot.Proxies[proxyID]; !desired {
continue
}
agent, ok := p.active.getByProxyID(proxyID)
if !ok {
continue
}
// Wait until all desired agents have a measured RTT before
// choosing any to disconnect.
rtt, ok := agent.RTT()
if !ok {
return
}
if rtt > maxRTT {
maxRTT = rtt
disconnectCandidate = agent
}
}
if disconnectCandidate != nil {
p.lastConnectivityChange = now
go func() {
if err := disconnectCandidate.Stop(); err != nil {
p.logger.DebugContext(ctx, "Error disconnecting agent", "error", err)
}
}()
}
}
// waitForBackoff processes events while waiting for the backoff.
func (p *AgentPool) waitForBackoff(ctx context.Context, events <-chan Agent) error {
for {
@@ -415,12 +502,20 @@ func (p *AgentPool) waitForBackoff(ctx context.Context, events <-chan Agent) err
p.backoff.Inc()
return nil
case agent := <-events:
p.handleEvent(ctx, agent)
p.handleSyncEvent(ctx, agent)
}
}
}
// handleEvent processes a single event.
// handleSyncEvent handles a single agent event. This should be used synchonously
// within the main agentpool loop.
func (p *AgentPool) handleSyncEvent(ctx context.Context, agent Agent) {
p.lastConnectivityChange = p.Clock.Now()
p.handleEvent(ctx, agent)
}
// handleEvent processes a single agent event. Use [AgentPool.handleSyncEvent]
// instead when executing from within the main agentpool loop.
func (p *AgentPool) handleEvent(ctx context.Context, agent Agent) {
state := agent.GetState()
switch state {
@@ -647,6 +742,9 @@ type agentPoolRuntimeConfig struct {
// connectionCount determines how many proxy servers the agent pool will
// connect to. This settings is ignored for the AgentMesh tunnel strategy.
connectionCount int
// disconnectThreshold is the minimum stable period before excess proxy
// peering connections may be disconnected.
disconnectThreshold time.Duration
// keepAliveInterval is the interval agents will send heartbeats at.
keepAliveInterval time.Duration
// keepAliveCount specifies the amount of missed ping heartbeats
@@ -712,6 +810,12 @@ func (c *agentPoolRuntimeConfig) getConnectionCount() int {
return c.connectionCount
}
func (c *agentPoolRuntimeConfig) getDisconnectThreshold() time.Duration {
c.mu.RLock()
defer c.mu.RUnlock()
return c.disconnectThreshold
}
// useReverseTunnelV2Locked returns true if reverse tunnel should be used.
func (c *agentPoolRuntimeConfig) useReverseTunnelV2Locked() bool {
if c.isRemoteCluster {
@@ -836,6 +940,7 @@ func (c *agentPoolRuntimeConfig) update(ctx context.Context, netConfig types.Clu
if c.tunnelStrategyType == types.ProxyPeering {
strategy := netConfig.GetProxyPeeringTunnelStrategy()
c.connectionCount = int(strategy.AgentConnectionCount)
c.disconnectThreshold = time.Duration(strategy.DisconnectThresholdSeconds) * time.Second
}
if c.connectionCount <= 0 {
c.connectionCount = defaultAgentConnectionCount
+204 -1
View File
@@ -42,7 +42,11 @@ import (
type mockAgent struct {
Agent
proxyID string
rtt time.Duration
stopCalls int
mockStart func(ctx context.Context) error
mockStop func() error
mockGetState func() AgentState
}
@@ -61,7 +65,25 @@ func (m *mockAgent) GetState() AgentState {
}
func (m *mockAgent) GetProxyID() (string, bool) {
return "test-id", true
if m.proxyID == "" {
return "", false
}
return m.proxyID, true
}
func (m *mockAgent) Stop() error {
m.stopCalls++
if m.mockStop != nil {
return m.mockStop()
}
return nil
}
func (m *mockAgent) RTT() (time.Duration, bool) {
if m.rtt == 0 {
return 0, false
}
return m.rtt, true
}
type mockClient struct {
@@ -179,6 +201,187 @@ type mockSigner struct {
ssh.Signer
}
func TestAgentPoolTryDisconnect(t *testing.T) {
const disconnectThreshold = time.Minute
newTestPool := func(t *testing.T) *AgentPool {
client := &mockClient{}
pool, err := NewAgentPool(context.Background(), AgentPoolConfig{
Client: client,
AccessPoint: client,
PublicKeyAuth: apissh.PublicKeyAuthConfig{
Signers: func() ([]ssh.Signer, error) {
return []ssh.Signer{mockSigner{}}, nil
},
},
HostUUID: "test-uuid",
LocalCluster: "test-cluster",
Cluster: "test-cluster",
Resolver: func(context.Context) (*utils.NetAddr, types.ProxyListenerMode, error) {
return &utils.NetAddr{}, types.ProxyListenerMode_Separate, nil
},
})
require.NoError(t, err)
pool.runtimeConfig.disconnectThreshold = disconnectThreshold
pool.lastConnectivityChange = time.Now()
return pool
}
trackedProxiesForAgents := func(agents []*mockAgent) []track.Proxy {
proxies := make([]track.Proxy, 0, len(agents))
for _, agent := range agents {
if agent.proxyID == "" {
continue
}
proxies = append(proxies, track.Proxy{Name: agent.proxyID})
}
return proxies
}
type testUpdate struct {
apply func(*testing.T, *AgentPool)
assert func(*testing.T, *AgentPool, []*mockAgent)
}
tests := []struct {
name string
agents []*mockAgent
update []testUpdate
}{
{
name: "wait for all agents to report rtt",
agents: []*mockAgent{
{proxyID: "proxy-1", rtt: 10 * time.Millisecond},
{proxyID: "proxy-2"},
{proxyID: "proxy-3", rtt: 30 * time.Millisecond},
},
update: []testUpdate{
{
apply: func(t *testing.T, pool *AgentPool) {
pool.tracker.SetConnectionCount(2)
time.Sleep(disconnectThreshold + time.Second)
},
assert: func(t *testing.T, pool *AgentPool, agents []*mockAgent) {
require.Zero(t, agents[0].stopCalls)
require.Zero(t, agents[1].stopCalls)
require.Zero(t, agents[2].stopCalls)
require.NotEqual(t, time.Now(), pool.lastConnectivityChange)
},
},
},
},
{
name: "disconnect agent with highest rtt",
agents: []*mockAgent{
{proxyID: "proxy-1", rtt: 10 * time.Millisecond},
{proxyID: "proxy-2", rtt: 70 * time.Millisecond},
{proxyID: "proxy-3", rtt: 30 * time.Millisecond},
},
update: []testUpdate{
{
apply: func(t *testing.T, pool *AgentPool) {
pool.tracker.SetConnectionCount(2)
time.Sleep(disconnectThreshold + time.Second)
},
assert: func(t *testing.T, pool *AgentPool, agents []*mockAgent) {
require.Zero(t, agents[0].stopCalls)
require.Equal(t, 1, agents[1].stopCalls)
require.Zero(t, agents[2].stopCalls)
require.Equal(t, time.Now(), pool.lastConnectivityChange)
},
},
},
},
{
name: "wait for disconnect threshold after last connectivity change",
agents: []*mockAgent{
{proxyID: "proxy-1", rtt: 10 * time.Millisecond},
{proxyID: "proxy-2", rtt: 70 * time.Millisecond},
{proxyID: "proxy-3", rtt: 30 * time.Millisecond},
},
update: []testUpdate{
{
apply: func(t *testing.T, pool *AgentPool) {
pool.tracker.SetConnectionCount(2)
time.Sleep(disconnectThreshold + time.Second)
pool.lastConnectivityChange = time.Now()
},
assert: func(t *testing.T, pool *AgentPool, agents []*mockAgent) {
require.Zero(t, agents[0].stopCalls)
require.Zero(t, agents[1].stopCalls)
require.Zero(t, agents[2].stopCalls)
require.Equal(t, time.Now(), pool.lastConnectivityChange)
},
},
{
apply: func(t *testing.T, pool *AgentPool) {
time.Sleep(disconnectThreshold + time.Second)
},
assert: func(t *testing.T, pool *AgentPool, agents []*mockAgent) {
require.Zero(t, agents[0].stopCalls)
require.Equal(t, 1, agents[1].stopCalls)
require.Zero(t, agents[2].stopCalls)
require.Equal(t, time.Now(), pool.lastConnectivityChange)
},
},
},
},
{
name: "wait for disconnect threshold after last topology change",
agents: []*mockAgent{
{proxyID: "proxy-1", rtt: 10 * time.Millisecond},
{proxyID: "proxy-2", rtt: 70 * time.Millisecond},
{proxyID: "proxy-3", rtt: 30 * time.Millisecond},
},
update: []testUpdate{
{
apply: func(t *testing.T, pool *AgentPool) {
pool.tracker.SetConnectionCount(2)
time.Sleep(disconnectThreshold + time.Second)
pool.tracker.TrackExpected(track.Proxy{Name: "proxy-4"})
},
assert: func(t *testing.T, pool *AgentPool, agents []*mockAgent) {
require.Zero(t, agents[0].stopCalls)
require.Zero(t, agents[1].stopCalls)
require.Zero(t, agents[2].stopCalls)
require.NotEqual(t, time.Now(), pool.lastConnectivityChange)
require.Equal(t, time.Now(), pool.tracker.Snapshot().LastTopologyChange)
},
},
{
apply: func(t *testing.T, pool *AgentPool) {
time.Sleep(disconnectThreshold + time.Second)
},
assert: func(t *testing.T, pool *AgentPool, agents []*mockAgent) {
require.Zero(t, agents[0].stopCalls)
require.Equal(t, 1, agents[1].stopCalls)
require.Zero(t, agents[2].stopCalls)
require.NotEqual(t, time.Now(), pool.tracker.Snapshot().LastTopologyChange)
require.Equal(t, time.Now(), pool.lastConnectivityChange)
},
},
},
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
synctest.Test(t, func(t *testing.T) {
pool := newTestPool(t)
for _, agent := range tt.agents {
pool.active.add(agent)
}
pool.tracker.TrackExpected(trackedProxiesForAgents(tt.agents)...)
for _, update := range tt.update {
update.apply(t, pool)
pool.tryDisconnect(t.Context())
synctest.Wait()
update.assert(t, pool, tt.agents)
}
})
})
}
}
func TestAgentKeepAliveCountTimeout(t *testing.T) {
cases := []struct {
name string
+61 -6
View File
@@ -62,6 +62,10 @@ type Tracker struct {
mu sync.Mutex
// LastTopologyChange is the last time a state change was observed. State
// changes include: proxy membership, group/generation, and connection count.
lastTopologyChange time.Time
// connectionCount is nonpositive for full connectivity (agent mesh) mode, a
// positive number for the connection count of proxy peering mode.
connectionCount int
@@ -118,10 +122,11 @@ func New(cfg Config) (*Tracker, error) {
return nil, trace.Wrap(err)
}
t := &Tracker{
proxyExpiry: cfg.ProxyExpiry,
clusterSuffix: "." + cfg.ClusterName,
claimed: make(map[string]struct{}),
tracked: make(map[string]Proxy),
proxyExpiry: cfg.ProxyExpiry,
clusterSuffix: "." + cfg.ClusterName,
lastTopologyChange: time.Now(),
claimed: make(map[string]struct{}),
tracked: make(map[string]Proxy),
}
return t, nil
}
@@ -163,6 +168,7 @@ func (t *Tracker) expireProxiesLocked() {
if v.expiry.Before(now) {
delete(t.tracked, k)
t.cannotLease = false
t.lastTopologyChange = now
}
}
}
@@ -218,11 +224,20 @@ func (t *Tracker) TrackExpected(proxies ...Proxy) {
t.mu.Lock()
defer t.mu.Unlock()
now := time.Now()
expiry := time.Now().Add(t.proxyExpiry)
t.cannotLease = false
for _, p := range proxies {
p.expiry = expiry
ttl := p.TTL
if ttl <= 0 {
ttl = t.proxyExpiry
}
prev, ok := t.tracked[p.Name]
if !ok || prev.Group != p.Group || prev.Generation != p.Generation {
t.lastTopologyChange = now
}
p.TTL = ttl
p.expiry = now.Add(ttl)
// TODO(espadolini): log a warning if a proxy is changing group or
// changing (decreasing?) generation; it should be paired with support
// for parsing discovery messages coming from proxies that don't support
@@ -239,11 +254,51 @@ func (t *Tracker) TrackExpected(proxies ...Proxy) {
func (t *Tracker) SetConnectionCount(connectionCount int) {
t.mu.Lock()
defer t.mu.Unlock()
now := time.Now()
if t.connectionCount != connectionCount {
t.lastTopologyChange = now
}
t.cannotLease = false
t.connectionCount = connectionCount
}
// TrackerSnapshot is a snapshot of the trackers state.
type TrackerSnapshot struct {
// Proxies is a map of proxy ID to a bool indicating if the proxy is in the
// desired set. A desired proxy is in the latest generation of a group.
Proxies map[string]bool
// ConnectionCount is the configured connection count used to issue leases.
ConnectionCount int
// LastTopologyChange is the last time a state change was observed. State
// changes include: proxy membership, group/generation, and connection count.
LastTopologyChange time.Time
}
// Snapshot returns a snapshot of the tracker state.
func (t *Tracker) Snapshot() TrackerSnapshot {
t.mu.Lock()
defer t.mu.Unlock()
snapshot := TrackerSnapshot{
Proxies: make(map[string]bool, len(t.tracked)),
ConnectionCount: t.connectionCount,
LastTopologyChange: t.lastTopologyChange,
}
desiredGen := make(map[string]uint64, 8)
for _, v := range t.tracked {
if v.Generation > desiredGen[v.Group] {
desiredGen[v.Group] = v.Generation
}
}
for k, v := range t.tracked {
snapshot.Proxies[k] = desiredGen[v.Group] == v.Generation
}
return snapshot
}
// Claim attempts to claim exclusive access to a reverse tunnel server
// identified by the principals. It will fail if the server is already claimed
// or if the Lease had already been released or has already claimed a different
+72 -25
View File
@@ -23,6 +23,7 @@ import (
"math/rand/v2"
"sync"
"testing"
"testing/synctest"
"time"
"github.com/stretchr/testify/require"
@@ -134,19 +135,24 @@ func jitter(t time.Duration) time.Duration {
return t + rand.N(t/5)
}
func newTestTracker(t *testing.T) *Tracker {
tracker, err := New(Config{
ClusterName: "test-cluster",
})
require.NoError(t, err)
return tracker
}
func TestBasic(t *testing.T) {
const (
timeout = time.Second * 16
proxyCount = 16
)
tracker := newTestTracker(t)
timeoutC := time.After(timeout)
ticker := time.NewTicker(time.Millisecond * 10)
t.Cleanup(ticker.Stop)
tracker, err := New(Config{ClusterName: "test-cluster"})
require.NoError(t, err)
min, max := time.Duration(0), timeout
var proxies simpleTestProxies
proxies := simpleTestProxies{}
proxies.AddRandProxies(proxyCount, min, max)
for {
@@ -162,10 +168,11 @@ func TestBasic(t *testing.T) {
}
select {
case <-ticker.C:
t.Logf("activeCount: %v", tracker.activeCount())
case <-timeoutC:
t.Fatal("timeout")
default:
time.Sleep(10 * time.Millisecond)
t.Logf("activeCount: %v", tracker.activeCount())
}
}
}
@@ -180,13 +187,9 @@ func TestFullRotation(t *testing.T) {
timeout = time.Second * 30
)
ticker := time.NewTicker(time.Millisecond * 100)
t.Cleanup(ticker.Stop)
var proxies simpleTestProxies
tracker := newTestTracker(t)
proxies := simpleTestProxies{}
proxies.AddRandProxies(proxyCount, minConnA, maxConnA)
tracker, err := New(Config{ClusterName: "test-cluster"})
require.NoError(t, err)
timeoutC := time.After(timeout)
for {
@@ -208,10 +211,11 @@ func TestFullRotation(t *testing.T) {
}
select {
case <-ticker.C:
t.Logf("activeCount0: %v", tracker.activeCount())
case <-timeoutC:
t.Fatal("timeout")
default:
time.Sleep(100 * time.Millisecond)
t.Logf("activeCount0: %v", tracker.activeCount())
}
}
proxies.RemoveRandProxies(proxyCount)
@@ -223,10 +227,11 @@ func TestFullRotation(t *testing.T) {
}
select {
case <-ticker.C:
t.Logf("activeCount1: %v", tracker.activeCount())
case <-timeoutC:
t.Fatal("timeout")
default:
time.Sleep(100 * time.Millisecond)
t.Logf("activeCount1: %v", tracker.activeCount())
}
}
proxies.AddRandProxies(proxyCount, minConnB, maxConnB)
@@ -243,10 +248,11 @@ func TestFullRotation(t *testing.T) {
}
select {
case <-ticker.C:
t.Logf("activeCount2: %v", tracker.activeCount())
case <-timeoutC:
t.Fatal("timeout")
default:
time.Sleep(100 * time.Millisecond)
t.Logf("activeCount2 %v", tracker.activeCount())
}
}
}
@@ -255,8 +261,7 @@ func TestFullRotation(t *testing.T) {
// from the expected teleport principal format, and that gossip messages
// consisting only of uuid don't create duplicate entries.
func TestUUIDHandling(t *testing.T) {
tracker, err := New(Config{ClusterName: "test-cluster"})
require.NoError(t, err)
tracker := newTestTracker(t)
lease := tracker.TryAcquire()
require.NotNil(t, lease)
@@ -269,8 +274,7 @@ func TestUUIDHandling(t *testing.T) {
}
func TestIsClaimed(t *testing.T) {
tracker, err := New(Config{ClusterName: "test-cluster"})
require.NoError(t, err)
tracker := newTestTracker(t)
tracker.TrackExpected(Proxy{Name: "proxy1"}, Proxy{Name: "proxy2"})
require.False(t, tracker.IsClaimed("proxy1.test-cluster"))
@@ -299,8 +303,7 @@ func (t *Tracker) activeCount() int {
}
func TestProxyGroups(t *testing.T) {
tracker, err := New(Config{ClusterName: "test-cluster"})
require.NoError(t, err)
tracker := newTestTracker(t)
tracker.SetConnectionCount(2)
@@ -361,3 +364,47 @@ func TestProxyGroups(t *testing.T) {
yc.Release()
requireAcquire()
}
func TestProxyTTL(t *testing.T) {
synctest.Test(t, func(t *testing.T) {
tracker, err := New(Config{
ClusterName: "test-cluster",
ProxyExpiry: time.Second,
})
require.NoError(t, err)
tracker.TrackExpected(Proxy{
Name: "proxy-1",
TTL: 100 * time.Millisecond,
})
lease := tracker.TryAcquire()
require.NotNil(t, lease)
require.True(t, lease.Claim("proxy-1"))
time.Sleep(200 * time.Millisecond)
require.NotNil(t, tracker.TryAcquire())
})
}
func TestDefaultProxyExpiry(t *testing.T) {
synctest.Test(t, func(t *testing.T) {
tracker, err := New(Config{
ClusterName: "test-cluster",
ProxyExpiry: time.Second,
})
require.NoError(t, err)
tracker.TrackExpected(Proxy{Name: "proxy-1"})
lease := tracker.TryAcquire()
require.NotNil(t, lease)
require.True(t, lease.Claim("proxy-1"))
time.Sleep(time.Second / 2)
require.Nil(t, tracker.TryAcquire())
time.Sleep(time.Second)
require.NotNil(t, tracker.TryAcquire())
})
}