mirror of
https://github.com/gravitational/teleport.git
synced 2026-10-12 02:15:01 +00:00
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:
@@ -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.
|
||||
|
||||
@@ -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
File diff suppressed because it is too large
Load Diff
+1
@@ -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.
|
||||
|
||||
|
||||
+1
@@ -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 {
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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
|
||||
}
|
||||
|
||||
@@ -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)
|
||||
}
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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())
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user