Skip to content

Commit 6e7e25f

Browse files
spikecurtisspikecurtis
committed
feat: use Agent v2 API for Service Banner
1 parent 6f243f6 commit 6e7e25f

6 files changed

Lines changed: 203 additions & 112 deletions

File tree

‎agent/agent.go‎

Lines changed: 31 additions & 20 deletions
Original file line numberDiff line numberDiff line change
@@ -41,6 +41,7 @@ import (
4141
"github.com/coder/coder/v2/agent/agentproc"
4242
"github.com/coder/coder/v2/agent/agentscripts"
4343
"github.com/coder/coder/v2/agent/agentssh"
44+
"github.com/coder/coder/v2/agent/proto"
4445
"github.com/coder/coder/v2/agent/reconnectingpty"
4546
"github.com/coder/coder/v2/buildinfo"
4647
"github.com/coder/coder/v2/cli/gitauth"
@@ -95,7 +96,6 @@ type Client interface {
9596
PostStartup(ctx context.Context, req agentsdk.PostStartupRequest) error
9697
PostMetadata(ctx context.Context, req agentsdk.PostMetadataRequest) error
9798
PatchLogs(ctx context.Context, req agentsdk.PatchLogs) error
98-
GetServiceBanner(ctx context.Context) (codersdk.ServiceBannerConfig, error)
9999
}
100100

101101
type Agent interface {
@@ -269,7 +269,6 @@ func (a *agent) init(ctx context.Context) {
269269
func (a *agent) runLoop(ctx context.Context) {
270270
go a.reportLifecycleLoop(ctx)
271271
go a.reportMetadataLoop(ctx)
272-
go a.fetchServiceBannerLoop(ctx)
273272
go a.manageProcessPriorityLoop(ctx)
274273

275274
for retrier := retry.New(100*time.Millisecond, 10*time.Second); retrier.Wait(ctx); {
@@ -662,22 +661,23 @@ func (a *agent) setLifecycle(ctx context.Context, state codersdk.WorkspaceAgentL
662661
// fetchServiceBannerLoop fetches the service banner on an interval. It will
663662
// not be fetched immediately; the expectation is that it is primed elsewhere
664663
// (and must be done before the session actually starts).
665-
func (a *agent) fetchServiceBannerLoop(ctx context.Context) {
664+
func (a *agent) fetchServiceBannerLoop(ctx context.Context, aAPI proto.DRPCAgentClient) error {
666665
ticker := time.NewTicker(a.serviceBannerRefreshInterval)
667666
defer ticker.Stop()
668667
for {
669668
select {
670669
case <-ctx.Done():
671-
return
670+
return ctx.Err()
672671
case <-ticker.C:
673-
serviceBanner, err := a.client.GetServiceBanner(ctx)
672+
sbp, err := aAPI.GetServiceBanner(ctx, &proto.GetServiceBannerRequest{})
674673
if err != nil {
675674
if ctx.Err() != nil {
676-
return
675+
return ctx.Err()
677676
}
678677
a.logger.Error(ctx, "failed to update service banner", slog.Error(err))
679-
continue
678+
return err
680679
}
680+
serviceBanner := proto.SDKServiceBannerFromProto(sbp)
681681
a.serviceBanner.Store(&serviceBanner)
682682
}
683683
}
@@ -693,10 +693,24 @@ func (a *agent) run(ctx context.Context) error {
693693
}
694694
a.sessionToken.Store(&sessionToken)
695695

696-
serviceBanner, err := a.client.GetServiceBanner(ctx)
696+
// Listen returns the dRPC connection we use for the Agent v2+ API
697+
conn, err := a.client.Listen(ctx)
698+
if err != nil {
699+
return err
700+
}
701+
defer func() {
702+
cErr := conn.Close()
703+
if cErr != nil {
704+
a.logger.Debug(ctx, "error closing drpc connection", slog.Error(err))
705+
}
706+
}()
707+
708+
aAPI := proto.NewDRPCAgentClient(conn)
709+
sbp, err := aAPI.GetServiceBanner(ctx, &proto.GetServiceBannerRequest{})
697710
if err != nil {
698711
return xerrors.Errorf("fetch service banner: %w", err)
699712
}
713+
serviceBanner := proto.SDKServiceBannerFromProto(sbp)
700714
a.serviceBanner.Store(&serviceBanner)
701715

702716
manifest, err := a.client.Manifest(ctx)
@@ -821,18 +835,6 @@ func (a *agent) run(ctx context.Context) error {
821835
network.SetBlockEndpoints(manifest.DisableDirectConnections)
822836
}
823837

824-
// Listen returns the dRPC connection we use for both Coordinator and DERPMap updates
825-
conn, err := a.client.Listen(ctx)
826-
if err != nil {
827-
return err
828-
}
829-
defer func() {
830-
cErr := conn.Close()
831-
if cErr != nil {
832-
a.logger.Debug(ctx, "error closing drpc connection", slog.Error(err))
833-
}
834-
}()
835-
836838
eg, egCtx := errgroup.WithContext(ctx)
837839
eg.Go(func() error {
838840
a.logger.Debug(egCtx, "running tailnet connection coordinator")
@@ -852,6 +854,15 @@ func (a *agent) run(ctx context.Context) error {
852854
return nil
853855
})
854856

857+
eg.Go(func() error {
858+
a.logger.Debug(egCtx, "running fetch server banner loop")
859+
err := a.fetchServiceBannerLoop(egCtx, aAPI)
860+
if err != nil {
861+
return xerrors.Errorf("fetch server banner loop: %w", err)
862+
}
863+
return nil
864+
})
865+
855866
return eg.Wait()
856867
}
857868

‎agent/agenttest/client.go‎

Lines changed: 87 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -18,6 +18,7 @@ import (
1818
"tailscale.com/tailcfg"
1919

2020
"cdr.dev/slog"
21+
agentproto "github.com/coder/coder/v2/agent/proto"
2122
"github.com/coder/coder/v2/codersdk"
2223
"github.com/coder/coder/v2/codersdk/agentsdk"
2324
drpcsdk "github.com/coder/coder/v2/codersdk/drpc"
@@ -48,6 +49,9 @@ func NewClient(t testing.TB,
4849
}
4950
err := proto.DRPCRegisterTailnet(mux, drpcService)
5051
require.NoError(t, err)
52+
fakeAAPI := NewFakeAgentAPI(t, logger)
53+
err = agentproto.DRPCRegisterAgent(mux, fakeAAPI)
54+
require.NoError(t, err)
5155
server := drpcserver.NewWithOptions(mux, drpcserver.Options{
5256
Log: func(err error) {
5357
if xerrors.Is(err, io.EOF) {
@@ -64,22 +68,23 @@ func NewClient(t testing.TB,
6468
statsChan: statsChan,
6569
coordinator: coordinator,
6670
server: server,
71+
fakeAgentAPI: fakeAAPI,
6772
derpMapUpdates: derpMapUpdates,
6873
}
6974
}
7075

7176
type Client struct {
72-
t testing.TB
73-
logger slog.Logger
74-
agentID uuid.UUID
75-
manifest agentsdk.Manifest
76-
metadata map[string]agentsdk.Metadata
77-
statsChan chan *agentsdk.Stats
78-
coordinator tailnet.Coordinator
79-
server *drpcserver.Server
80-
LastWorkspaceAgent func()
81-
PatchWorkspaceLogs func() error
82-
GetServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
77+
t testing.TB
78+
logger slog.Logger
79+
agentID uuid.UUID
80+
manifest agentsdk.Manifest
81+
metadata map[string]agentsdk.Metadata
82+
statsChan chan *agentsdk.Stats
83+
coordinator tailnet.Coordinator
84+
server *drpcserver.Server
85+
fakeAgentAPI *FakeAgentAPI
86+
LastWorkspaceAgent func()
87+
PatchWorkspaceLogs func() error
8388

8489
mu sync.Mutex // Protects following.
8590
lifecycleStates []codersdk.WorkspaceAgentLifecycle
@@ -221,20 +226,7 @@ func (c *Client) PatchLogs(ctx context.Context, logs agentsdk.PatchLogs) error {
221226
}
222227

223228
func (c *Client) SetServiceBannerFunc(f func() (codersdk.ServiceBannerConfig, error)) {
224-
c.mu.Lock()
225-
defer c.mu.Unlock()
226-
227-
c.GetServiceBannerFunc = f
228-
}
229-
230-
func (c *Client) GetServiceBanner(ctx context.Context) (codersdk.ServiceBannerConfig, error) {
231-
c.mu.Lock()
232-
defer c.mu.Unlock()
233-
c.logger.Debug(ctx, "get service banner")
234-
if c.GetServiceBannerFunc != nil {
235-
return c.GetServiceBannerFunc()
236-
}
237-
return codersdk.ServiceBannerConfig{}, nil
229+
c.fakeAgentAPI.SetServiceBannerFunc(f)
238230
}
239231

240232
func (c *Client) PushDERPMapUpdate(update *tailcfg.DERPMap) error {
@@ -254,3 +246,73 @@ type closeFunc func() error
254246
func (c closeFunc) Close() error {
255247
return c()
256248
}
249+
250+
type FakeAgentAPI struct {
251+
sync.Mutex
252+
t testing.TB
253+
logger slog.Logger
254+
255+
getServiceBannerFunc func() (codersdk.ServiceBannerConfig, error)
256+
}
257+
258+
func (*FakeAgentAPI) GetManifest(context.Context, *agentproto.GetManifestRequest) (*agentproto.Manifest, error) {
259+
// TODO implement me
260+
panic("implement me")
261+
}
262+
263+
func (f *FakeAgentAPI) SetServiceBannerFunc(fn func() (codersdk.ServiceBannerConfig, error)) {
264+
f.Lock()
265+
defer f.Unlock()
266+
f.getServiceBannerFunc = fn
267+
f.logger.Info(context.Background(), "updated ServiceBannerFunc")
268+
}
269+
270+
func (f *FakeAgentAPI) GetServiceBanner(context.Context, *agentproto.GetServiceBannerRequest) (*agentproto.ServiceBanner, error) {
271+
f.Lock()
272+
defer f.Unlock()
273+
if f.getServiceBannerFunc == nil {
274+
return &agentproto.ServiceBanner{}, nil
275+
}
276+
sb, err := f.getServiceBannerFunc()
277+
if err != nil {
278+
return nil, err
279+
}
280+
return agentproto.ServiceBannerFromSDK(sb), nil
281+
}
282+
283+
func (*FakeAgentAPI) UpdateStats(context.Context, *agentproto.UpdateStatsRequest) (*agentproto.UpdateStatsResponse, error) {
284+
// TODO implement me
285+
panic("implement me")
286+
}
287+
288+
func (*FakeAgentAPI) UpdateLifecycle(context.Context, *agentproto.UpdateLifecycleRequest) (*agentproto.Lifecycle, error) {
289+
// TODO implement me
290+
panic("implement me")
291+
}
292+
293+
func (*FakeAgentAPI) BatchUpdateAppHealths(context.Context, *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
294+
// TODO implement me
295+
panic("implement me")
296+
}
297+
298+
func (*FakeAgentAPI) UpdateStartup(context.Context, *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {
299+
// TODO implement me
300+
panic("implement me")
301+
}
302+
303+
func (*FakeAgentAPI) BatchUpdateMetadata(context.Context, *agentproto.BatchUpdateMetadataRequest) (*agentproto.BatchUpdateMetadataResponse, error) {
304+
// TODO implement me
305+
panic("implement me")
306+
}
307+
308+
func (*FakeAgentAPI) BatchCreateLogs(context.Context, *agentproto.BatchCreateLogsRequest) (*agentproto.BatchCreateLogsResponse, error) {
309+
// TODO implement me
310+
panic("implement me")
311+
}
312+
313+
func NewFakeAgentAPI(t testing.TB, logger slog.Logger) *FakeAgentAPI {
314+
return &FakeAgentAPI{
315+
t: t,
316+
logger: logger.Named("FakeAgentAPI"),
317+
}
318+
}

‎cli/agent_test.go‎

Lines changed: 33 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -212,7 +212,39 @@ func TestWorkspaceAgent(t *testing.T) {
212212

213213
clitest.Start(t, inv)
214214

215-
resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
215+
ctx := testutil.Context(t, testutil.WaitLong)
216+
// here we can't use coderdtest.AwaitWorkspaceAgents because it just waits
217+
// for the agent to connect, but we are testing PostStartup which can
218+
// happen after the Agent connects, so we just have to wait for it.
219+
var resources []codersdk.WorkspaceResource
220+
require.Eventually(t, func() bool {
221+
var err error
222+
workspace, err := client.Workspace(ctx, r.Workspace.ID)
223+
if !assert.NoError(t, err) {
224+
return false
225+
}
226+
if workspace.LatestBuild.Job.CompletedAt == nil {
227+
return false
228+
}
229+
if workspace.LatestBuild.Job.CompletedAt.IsZero() {
230+
return false
231+
}
232+
233+
resources = workspace.LatestBuild.Resources
234+
if len(resources) < 1 {
235+
t.Log("workspace has no resources")
236+
return false
237+
}
238+
if len(resources[0].Agents) < 1 {
239+
t.Log("workspace has no agents")
240+
return false
241+
}
242+
if len(resources[0].Agents[0].Subsystems) < 2 {
243+
t.Logf("not enough subsystems %v", resources[0].Agents[0].Subsystems)
244+
return false
245+
}
246+
return true
247+
}, testutil.WaitLong, testutil.IntervalMedium)
216248
require.Len(t, resources, 1)
217249
require.Len(t, resources[0].Agents, 1)
218250
require.Len(t, resources[0].Agents[0].Subsystems, 2)

0 commit comments

Comments
 (0)