test(agent): fix service banner and metadata intervals (#8516)

This commit is contained in:
Mathias Fredriksson
2023-07-14 16:10:26 +03:00
committed by GitHub
parent b7806fd216
commit 5fd77ad7cf
2 changed files with 73 additions and 73 deletions

View File

@ -6,7 +6,6 @@ import (
"encoding/binary"
"encoding/json"
"errors"
"flag"
"fmt"
"io"
"net"
@ -65,8 +64,9 @@ type Options struct {
TailnetListenPort uint16
Subsystem codersdk.AgentSubsystem
Addresses []netip.Prefix
PrometheusRegistry *prometheus.Registry
ReportMetadataInterval time.Duration
ServiceBannerRefreshInterval time.Duration
}
type Client interface {
@ -107,6 +107,12 @@ func New(options Options) Agent {
return "", nil
}
}
if options.ReportMetadataInterval == 0 {
options.ReportMetadataInterval = 1 * time.Minute
}
if options.ServiceBannerRefreshInterval == 0 {
options.ServiceBannerRefreshInterval = 2 * time.Minute
}
prometheusRegistry := options.PrometheusRegistry
if prometheusRegistry == nil {
@ -131,6 +137,8 @@ func New(options Options) Agent {
lifecycleStates: []agentsdk.PostLifecycleRequest{{State: codersdk.WorkspaceAgentLifecycleCreated}},
ignorePorts: options.IgnorePorts,
connStatsChan: make(chan *agentsdk.Stats, 1),
reportMetadataInterval: options.ReportMetadataInterval,
serviceBannerRefreshInterval: options.ServiceBannerRefreshInterval,
sshMaxTimeout: options.SSHMaxTimeout,
subsystem: options.Subsystem,
addresses: options.Addresses,
@ -165,10 +173,11 @@ type agent struct {
closed chan struct{}
envVars map[string]string
// manifest is atomic because values can change after reconnection.
manifest atomic.Pointer[agentsdk.Manifest]
// serviceBanner is atomic because it can change.
serviceBanner atomic.Pointer[codersdk.ServiceBannerConfig]
manifest atomic.Pointer[agentsdk.Manifest] // manifest is atomic because values can change after reconnection.
reportMetadataInterval time.Duration
serviceBanner atomic.Pointer[codersdk.ServiceBannerConfig] // serviceBanner is atomic because it is periodically updated.
serviceBannerRefreshInterval time.Duration
sessionToken atomic.Pointer[string]
sshServer *agentssh.Server
sshMaxTimeout time.Duration
@ -283,17 +292,6 @@ func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentM
return result
}
// adjustIntervalForTests returns a duration of testInterval milliseconds long
// for tests and interval seconds long otherwise.
func adjustIntervalForTests(interval time.Duration, testInterval time.Duration) time.Duration {
// In tests we want to set shorter intervals because engineers are
// impatient.
if flag.Lookup("test.v") != nil {
return testInterval
}
return interval
}
type metadataResultAndKey struct {
result *codersdk.WorkspaceAgentMetadataResult
key string
@ -315,12 +313,10 @@ func (t *trySingleflight) Do(key string, fn func()) {
}
func (a *agent) reportMetadataLoop(ctx context.Context) {
baseInterval := adjustIntervalForTests(time.Second, time.Millisecond*100)
const metadataLimit = 128
var (
baseTicker = time.NewTicker(baseInterval)
baseTicker = time.NewTicker(a.reportMetadataInterval)
lastCollectedAts = make(map[string]time.Time)
metadataResults = make(chan metadataResultAndKey, metadataLimit)
)
@ -391,11 +387,7 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
continue
}
// The last collected value isn't quite stale yet, so we skip it.
if collectedAt.Add(
adjustIntervalForTests(
time.Duration(md.Interval)*time.Second,
time.Duration(md.Interval)*time.Millisecond*100),
).After(time.Now()) {
if collectedAt.Add(a.reportMetadataInterval).After(time.Now()) {
continue
}
}
@ -506,7 +498,7 @@ func (a *agent) setLifecycle(ctx context.Context, state codersdk.WorkspaceAgentL
// not be fetched immediately; the expectation is that it is primed elsewhere
// (and must be done before the session actually starts).
func (a *agent) fetchServiceBannerLoop(ctx context.Context) {
ticker := time.NewTicker(adjustIntervalForTests(2*time.Minute, time.Millisecond*5))
ticker := time.NewTicker(a.serviceBannerRefreshInterval)
defer ticker.Stop()
for {
select {

View File

@ -424,8 +424,12 @@ func TestAgent_Session_TTY_MOTD_Update(t *testing.T) {
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
setSBInterval := func(_ *agenttest.Client, opts *agent.Options) {
opts.ServiceBannerRefreshInterval = 5 * time.Millisecond
}
//nolint:dogsled // Allow the blank identifiers.
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0)
conn, client, _, _, _ := setupAgent(t, agentsdk.Manifest{}, 0, setSBInterval)
for _, test := range tests {
test := test
// 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) {
Script: echoHello,
},
},
}, 0)
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ReportMetadataInterval = 100 * time.Millisecond
})
var gotMd map[string]agentsdk.PostMetadataRequest
require.Eventually(t, func() bool {
@ -1174,23 +1180,23 @@ func TestAgent_Metadata(t *testing.T) {
Script: echoHello,
},
},
}, 0)
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ReportMetadataInterval = testutil.IntervalFast
})
var gotMd map[string]agentsdk.PostMetadataRequest
require.Eventually(t, func() bool {
gotMd = client.GetMetadata()
return len(gotMd) == 1
}, testutil.WaitShort, testutil.IntervalMedium)
}, testutil.WaitShort, testutil.IntervalFast/2)
collectedAt1 := gotMd["greeting"].CollectedAt
if !assert.Equal(t, "hello", strings.TrimSpace(gotMd["greeting"].Value)) {
t.Errorf("got: %+v", gotMd)
}
require.Equal(t, "hello", strings.TrimSpace(gotMd["greeting"].Value))
if !assert.Eventually(t, func() bool {
gotMd = client.GetMetadata()
return gotMd["greeting"].CollectedAt.After(collectedAt1)
}, testutil.WaitShort, testutil.IntervalMedium) {
}, testutil.WaitShort, testutil.IntervalFast/2) {
t.Fatalf("expected metadata to be collected again")
}
})
@ -1227,7 +1233,9 @@ func TestAgentMetadata_Timing(t *testing.T) {
Script: "exit 1",
},
},
}, 0)
}, 0, func(_ *agenttest.Client, opts *agent.Options) {
opts.ReportMetadataInterval = intervalUnit
})
require.Eventually(t, func() bool {
return len(client.GetMetadata()) == 2
@ -1844,14 +1852,14 @@ func setupSSHCommand(t *testing.T, beforeArgs []string, afterArgs []string) (*pt
func setupSSHSession(
t *testing.T,
options agentsdk.Manifest,
manifest agentsdk.Manifest,
serviceBanner codersdk.ServiceBannerConfig,
prepareFS func(fs afero.Fs),
) *ssh.Session {
ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitLong)
defer cancel()
//nolint:dogsled
conn, _, _, fs, _ := setupAgent(t, options, 0, func(c *agenttest.Client, _ *agent.Options) {
conn, _, _, fs, _ := setupAgent(t, manifest, 0, func(c *agenttest.Client, _ *agent.Options) {
c.SetServiceBannerFunc(func() (codersdk.ServiceBannerConfig, error) {
return serviceBanner, nil
})