Skip to content

Commit 5fd77ad

Browse files
mafredrimafredri
authored
test(agent): fix service banner and metadata intervals (coder#8516)
1 parent b7806fd commit 5fd77ad

File tree

2 files changed

+73
-73
lines changed

2 files changed

+73
-73
lines changed

‎agent/agent.go

+54-62
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,6 @@ import (
66
"encoding/binary"
77
"encoding/json"
88
"errors"
9-
"flag"
109
"fmt"
1110
"io"
1211
"net"
@@ -52,21 +51,22 @@ const (
5251
)
5352

5453
type Options struct {
55-
Filesystem afero.Fs
56-
LogDir string
57-
TempDir string
58-
ExchangeToken func(ctx context.Context) (string, error)
59-
Client Client
60-
ReconnectingPTYTimeout time.Duration
61-
EnvironmentVariables map[string]string
62-
Logger slog.Logger
63-
IgnorePorts map[int]string
64-
SSHMaxTimeout time.Duration
65-
TailnetListenPort uint16
66-
Subsystem codersdk.AgentSubsystem
67-
Addresses []netip.Prefix
68-
69-
PrometheusRegistry *prometheus.Registry
54+
Filesystem afero.Fs
55+
LogDir string
56+
TempDir string
57+
ExchangeToken func(ctx context.Context) (string, error)
58+
Client Client
59+
ReconnectingPTYTimeout time.Duration
60+
EnvironmentVariables map[string]string
61+
Logger slog.Logger
62+
IgnorePorts map[int]string
63+
SSHMaxTimeout time.Duration
64+
TailnetListenPort uint16
65+
Subsystem codersdk.AgentSubsystem
66+
Addresses []netip.Prefix
67+
PrometheusRegistry *prometheus.Registry
68+
ReportMetadataInterval time.Duration
69+
ServiceBannerRefreshInterval time.Duration
7070
}
7171

7272
type Client interface {
@@ -107,6 +107,12 @@ func New(options Options) Agent {
107107
return "", nil
108108
}
109109
}
110+
if options.ReportMetadataInterval == 0 {
111+
options.ReportMetadataInterval = 1 * time.Minute
112+
}
113+
if options.ServiceBannerRefreshInterval == 0 {
114+
options.ServiceBannerRefreshInterval = 2 * time.Minute
115+
}
110116

111117
prometheusRegistry := options.PrometheusRegistry
112118
if prometheusRegistry == nil {
@@ -115,25 +121,27 @@ func New(options Options) Agent {
115121

116122
ctx, cancelFunc := context.WithCancel(context.Background())
117123
a := &agent{
118-
tailnetListenPort: options.TailnetListenPort,
119-
reconnectingPTYTimeout: options.ReconnectingPTYTimeout,
120-
logger: options.Logger,
121-
closeCancel: cancelFunc,
122-
closed: make(chan struct{}),
123-
envVars: options.EnvironmentVariables,
124-
client: options.Client,
125-
exchangeToken: options.ExchangeToken,
126-
filesystem: options.Filesystem,
127-
logDir: options.LogDir,
128-
tempDir: options.TempDir,
129-
lifecycleUpdate: make(chan struct{}, 1),
130-
lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),
131-
lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
132-
ignorePorts: options.IgnorePorts,
133-
connStatsChan: make(chan *agentsdk.Stats, 1),
134-
sshMaxTimeout: options.SSHMaxTimeout,
135-
subsystem: options.Subsystem,
136-
addresses: options.Addresses,
124+
tailnetListenPort: options.TailnetListenPort,
125+
reconnectingPTYTimeout: options.ReconnectingPTYTimeout,
126+
logger: options.Logger,
127+
closeCancel: cancelFunc,
128+
closed: make(chan struct{}),
129+
envVars: options.EnvironmentVariables,
130+
client: options.Client,
131+
exchangeToken: options.ExchangeToken,
132+
filesystem: options.Filesystem,
133+
logDir: options.LogDir,
134+
tempDir: options.TempDir,
135+
lifecycleUpdate: make(chan struct{}, 1),
136+
lifecycleReported: make(chan codersdk.WorkspaceAgentLifecycle, 1),
137+
lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
138+
ignorePorts: options.IgnorePorts,
139+
connStatsChan: make(chan *agentsdk.Stats, 1),
140+
reportMetadataInterval: options.ReportMetadataInterval,
141+
serviceBannerRefreshInterval: options.ServiceBannerRefreshInterval,
142+
sshMaxTimeout: options.SSHMaxTimeout,
143+
subsystem: options.Subsystem,
144+
addresses: options.Addresses,
137145

138146
prometheusRegistry: prometheusRegistry,
139147
metrics: newAgentMetrics(prometheusRegistry),
@@ -165,13 +173,14 @@ type agent struct {
165173
closed chan struct{}
166174

167175
envVars map[string]string
168-
// manifest is atomic because values can change after reconnection.
169-
manifest atomic.Pointer[agentsdk.Manifest]
170-
// serviceBanner is atomic because it can change.
171-
serviceBanner atomic.Pointer[codersdk.ServiceBannerConfig]
172-
sessionToken atomic.Pointer[string]
173-
sshServer *agentssh.Server
174-
sshMaxTimeout time.Duration
176+
177+
manifest atomic.Pointer[agentsdk.Manifest] // manifest is atomic because values can change after reconnection.
178+
reportMetadataInterval time.Duration
179+
serviceBanner atomic.Pointer[codersdk.ServiceBannerConfig] // serviceBanner is atomic because it is periodically updated.
180+
serviceBannerRefreshInterval time.Duration
181+
sessionToken atomic.Pointer[string]
182+
sshServer *agentssh.Server
183+
sshMaxTimeout time.Duration
175184

176185
lifecycleUpdate chan struct{}
177186
lifecycleReported chan codersdk.WorkspaceAgentLifecycle
@@ -283,17 +292,6 @@ func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentM
283292
return result
284293
}
285294

286-
// adjustIntervalForTests returns a duration of testInterval milliseconds long
287-
// for tests and interval seconds long otherwise.
288-
func adjustIntervalForTests(interval time.Duration, testInterval time.Duration) time.Duration {
289-
// In tests we want to set shorter intervals because engineers are
290-
// impatient.
291-
if flag.Lookup("test.v") != nil {
292-
return testInterval
293-
}
294-
return interval
295-
}
296-
297295
type metadataResultAndKey struct {
298296
result *codersdk.WorkspaceAgentMetadataResult
299297
key string
@@ -315,12 +313,10 @@ func (t *trySingleflight) Do(key string, fn func()) {
315313
}
316314

317315
func (a *agent) reportMetadataLoop(ctx context.Context) {
318-
baseInterval := adjustIntervalForTests(time.Second, time.Millisecond*100)
319-
320316
const metadataLimit = 128
321317

322318
var (
323-
baseTicker = time.NewTicker(baseInterval)
319+
baseTicker = time.NewTicker(a.reportMetadataInterval)
324320
lastCollectedAts = make(map[string]time.Time)
325321
metadataResults = make(chan metadataResultAndKey, metadataLimit)
326322
)
@@ -391,11 +387,7 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
391387
continue
392388
}
393389
// The last collected value isn't quite stale yet, so we skip it.
394-
if collectedAt.Add(
395-
adjustIntervalForTests(
396-
time.Duration(md.Interval)*time.Second,
397-
time.Duration(md.Interval)*time.Millisecond*100),
398-
).After(time.Now()) {
390+
if collectedAt.Add(a.reportMetadataInterval).After(time.Now()) {
399391
continue
400392
}
401393
}
@@ -506,7 +498,7 @@ func (a *agent) setLifecycle(ctx context.Context, state codersdk.WorkspaceAgentL
506498
// not be fetched immediately; the expectation is that it is primed elsewhere
507499
// (and must be done before the session actually starts).
508500
func (a *agent) fetchServiceBannerLoop(ctx context.Context) {
509-
ticker := time.NewTicker(adjustIntervalForTests(2*time.Minute, time.Millisecond*5))
501+
ticker := time.NewTicker(a.serviceBannerRefreshInterval)
510502
defer ticker.Stop()
511503
for {
512504
select {

‎agent/agent_test.go

+19-11
Original file line numberDiff line numberDiff line change
@@ -424,8 +424,12 @@ func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
424424

425425
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
426426
defer cancel()
427+
428+
setSBInterval := func(_ *agenttest.Client, opts *agent.Options) {
429+
opts.ServiceBannerRefreshInterval = 5 * time.Millisecond
430+
}
427431
//nolint:dogsled // Allow the blank identifiers.
428-
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
432+
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, setSBInterval)
429433
for _, test := range tests {
430434
test := test
431435
// Set new banner func and wait for the agent to call it to update the
@@ -1143,7 +1147,9 @@ func TestAgent_Metadata(t *testing.T) {
11431147
Script: echoHello,
11441148
},
11451149
},
1146-
}, 0)
1150+
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
1151+
opts.ReportMetadataInterval = 100 * time.Millisecond
1152+
})
11471153

11481154
var gotMd map[string]agentsdk.PostMetadataRequest
11491155
require.Eventually(t, func() bool {
@@ -1174,23 +1180,23 @@ func TestAgent_Metadata(t *testing.T) {
11741180
Script: echoHello,
11751181
},
11761182
},
1177-
}, 0)
1183+
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
1184+
opts.ReportMetadataInterval = testutil.IntervalFast
1185+
})
11781186

11791187
var gotMd map[string]agentsdk.PostMetadataRequest
11801188
require.Eventually(t, func() bool {
11811189
gotMd = client.GetMetadata()
11821190
return len(gotMd) == 1
1183-
}, testutil.WaitShort, testutil.IntervalMedium)
1191+
}, testutil.WaitShort, testutil.IntervalFast/2)
11841192

11851193
collectedAt1 := gotMd["greeting"].CollectedAt
1186-
if !assert.Equal(t, "hello", strings.TrimSpace(gotMd["greeting"].Value)) {
1187-
t.Errorf("got: %+v", gotMd)
1188-
}
1194+
require.Equal(t, "hello", strings.TrimSpace(gotMd["greeting"].Value))
11891195

11901196
if !assert.Eventually(t, func() bool {
11911197
gotMd = client.GetMetadata()
11921198
return gotMd["greeting"].CollectedAt.After(collectedAt1)
1193-
}, testutil.WaitShort, testutil.IntervalMedium) {
1199+
}, testutil.WaitShort, testutil.IntervalFast/2) {
11941200
t.Fatalf("expected metadata to be collected again")
11951201
}
11961202
})
@@ -1227,7 +1233,9 @@ func TestAgentMetadata_Timing(t *testing.T) {
12271233
Script: "exit 1",
12281234
},
12291235
},
1230-
}, 0)
1236+
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
1237+
opts.ReportMetadataInterval = intervalUnit
1238+
})
12311239

12321240
require.Eventually(t, func() bool {
12331241
return len(client.GetMetadata()) == 2
@@ -1844,14 +1852,14 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) (*pt
18441852

18451853
func setupSSHSession(
18461854
t *testing.T,
1847-
options agentsdk.Manifest,
1855+
manifest agentsdk.Manifest,
18481856
serviceBanner codersdk.ServiceBannerConfig,
18491857
prepareFS func(fs afero.Fs),
18501858
) *ssh.Session {
18511859
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
18521860
defer cancel()
18531861
//nolint:dogsled
1854-
conn, _, _, fs, _ := setupAgent(t, options, 0, func(c *agenttest.Client, _ *agent.Options) {
1862+
conn, _, _, fs, _ := setupAgent(t, manifest, 0, func(c *agenttest.Client, _ *agent.Options) {
18551863
c.SetServiceBannerFunc(func() (codersdk.ServiceBannerConfig, error) {
18561864
return serviceBanner, nil
18571865
})

0 commit comments

Comments
 (0)