@@ -77,6 +77,7 @@ type Client interface {
77
77
PostStartup (ctx context.Context , req agentsdk.PostStartupRequest ) error
78
78
PostMetadata (ctx context.Context , key string , req agentsdk.PostMetadataRequest ) error
79
79
PatchStartupLogs (ctx context.Context , req agentsdk.PatchStartupLogs ) error
80
+ GetServiceBanner (ctx context.Context ) (codersdk.ServiceBannerConfig , error )
80
81
}
81
82
82
83
type Agent interface {
@@ -163,7 +164,9 @@ type agent struct {
163
164
164
165
envVars map [string ]string
165
166
// manifest is atomic because values can change after reconnection.
166
- manifest atomic.Pointer [agentsdk.Manifest ]
167
+ manifest atomic.Pointer [agentsdk.Manifest ]
168
+ // serviceBanner is atomic because it can change.
169
+ serviceBanner atomic.Pointer [codersdk.ServiceBannerConfig ]
167
170
sessionToken atomic.Pointer [string ]
168
171
sshServer * agentssh.Server
169
172
sshMaxTimeout time.Duration
@@ -191,6 +194,7 @@ func (a *agent) init(ctx context.Context) {
191
194
sshSrv .Env = a .envVars
192
195
sshSrv .AgentToken = func () string { return * a .sessionToken .Load () }
193
196
sshSrv .Manifest = & a .manifest
197
+ sshSrv .ServiceBanner = & a .serviceBanner
194
198
a .sshServer = sshSrv
195
199
196
200
go a .runLoop (ctx )
@@ -203,6 +207,7 @@ func (a *agent) init(ctx context.Context) {
203
207
func (a * agent ) runLoop (ctx context.Context ) {
204
208
go a .reportLifecycleLoop (ctx )
205
209
go a .reportMetadataLoop (ctx )
210
+ go a .fetchServiceBannerLoop (ctx )
206
211
207
212
for retrier := retry .New (100 * time .Millisecond , 10 * time .Second ); retrier .Wait (ctx ); {
208
213
a .logger .Info (ctx , "connecting to coderd" )
@@ -275,14 +280,15 @@ func (a *agent) collectMetadata(ctx context.Context, md codersdk.WorkspaceAgentM
275
280
return result
276
281
}
277
282
278
- func adjustIntervalForTests (i int64 ) time.Duration {
283
+ // adjustIntervalForTests returns a duration of testInterval milliseconds long
284
+ // for tests and interval seconds long otherwise.
285
+ func adjustIntervalForTests (interval time.Duration , testInterval time.Duration ) time.Duration {
279
286
// In tests we want to set shorter intervals because engineers are
280
287
// impatient.
281
- base := time .Second
282
288
if flag .Lookup ("test.v" ) != nil {
283
- base = time . Millisecond * 100
289
+ return testInterval
284
290
}
285
- return time . Duration ( i ) * base
291
+ return interval
286
292
}
287
293
288
294
type metadataResultAndKey struct {
@@ -306,7 +312,7 @@ func (t *trySingleflight) Do(key string, fn func()) {
306
312
}
307
313
308
314
func (a * agent ) reportMetadataLoop (ctx context.Context ) {
309
- baseInterval := adjustIntervalForTests (1 )
315
+ baseInterval := adjustIntervalForTests (time . Second , time . Millisecond * 100 )
310
316
311
317
const metadataLimit = 128
312
318
@@ -383,7 +389,9 @@ func (a *agent) reportMetadataLoop(ctx context.Context) {
383
389
}
384
390
// The last collected value isn't quite stale yet, so we skip it.
385
391
if collectedAt .Add (
386
- adjustIntervalForTests (md .Interval ),
392
+ adjustIntervalForTests (
393
+ time .Duration (md .Interval )* time .Second ,
394
+ time .Duration (md .Interval )* time .Millisecond * 100 ),
387
395
).After (time .Now ()) {
388
396
continue
389
397
}
@@ -491,6 +499,30 @@ func (a *agent) setLifecycle(ctx context.Context, state codersdk.WorkspaceAgentL
491
499
}
492
500
}
493
501
502
+ // fetchServiceBannerLoop fetches the service banner on an interval. It will
503
+ // not be fetched immediately; the expectation is that it is primed elsewhere
504
+ // (and must be done before the session actually starts).
505
+ func (a * agent ) fetchServiceBannerLoop (ctx context.Context ) {
506
+ ticker := time .NewTicker (adjustIntervalForTests (2 * time .Minute , time .Millisecond * 100 ))
507
+ defer ticker .Stop ()
508
+ for {
509
+ select {
510
+ case <- ctx .Done ():
511
+ return
512
+ case <- ticker .C :
513
+ serviceBanner , err := a .client .GetServiceBanner (ctx )
514
+ if err != nil {
515
+ if ctx .Err () != nil {
516
+ return
517
+ }
518
+ a .logger .Error (ctx , "failed to update service banner" , slog .Error (err ))
519
+ continue
520
+ }
521
+ a .serviceBanner .Store (& serviceBanner )
522
+ }
523
+ }
524
+ }
525
+
494
526
func (a * agent ) run (ctx context.Context ) error {
495
527
// This allows the agent to refresh it's token if necessary.
496
528
// For instance identity this is required, since the instance
@@ -501,6 +533,12 @@ func (a *agent) run(ctx context.Context) error {
501
533
}
502
534
a .sessionToken .Store (& sessionToken )
503
535
536
+ serviceBanner , err := a .client .GetServiceBanner (ctx )
537
+ if err != nil {
538
+ return xerrors .Errorf ("fetch service banner: %w" , err )
539
+ }
540
+ a .serviceBanner .Store (& serviceBanner )
541
+
504
542
manifest , err := a .client .Manifest (ctx )
505
543
if err != nil {
506
544
return xerrors .Errorf ("fetch metadata: %w" , err )
0 commit comments