Skip to content

Commit 320f232

Browse files
spikecurtisspikecurtis
committed
feat: use Agent v2 API for Service Banner
1 parent 385e74a commit 320f232

File tree

8 files changed

+265
-128
lines changed

8 files changed

+265
-128
lines changed

‎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: 36 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -83,7 +83,8 @@ func TestWorkspaceAgent(t *testing.T) {
8383

8484
ctx := inv.Context()
8585
clitest.Start(t, inv)
86-
coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
86+
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
87+
MatchResources(matchAgentWithVersion).Wait()
8788
workspace, err := client.Workspace(ctx, r.Workspace.ID)
8889
require.NoError(t, err)
8990
resources := workspace.LatestBuild.Resources
@@ -120,7 +121,9 @@ func TestWorkspaceAgent(t *testing.T) {
120121

121122
clitest.Start(t, inv)
122123
ctx := inv.Context()
123-
coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
124+
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
125+
MatchResources(matchAgentWithVersion).
126+
Wait()
124127
workspace, err := client.Workspace(ctx, r.Workspace.ID)
125128
require.NoError(t, err)
126129
resources := workspace.LatestBuild.Resources
@@ -161,7 +164,9 @@ func TestWorkspaceAgent(t *testing.T) {
161164
)
162165

163166
ctx := inv.Context()
164-
coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
167+
coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
168+
MatchResources(matchAgentWithVersion).
169+
Wait()
165170
workspace, err := client.Workspace(ctx, r.Workspace.ID)
166171
require.NoError(t, err)
167172
resources := workspace.LatestBuild.Resources
@@ -212,7 +217,8 @@ func TestWorkspaceAgent(t *testing.T) {
212217

213218
clitest.Start(t, inv)
214219

215-
resources := coderdtest.AwaitWorkspaceAgents(t, client, r.Workspace.ID)
220+
resources := coderdtest.NewWorkspaceAgentWaiter(t, client, r.Workspace.ID).
221+
MatchResources(matchAgentWithSubsystems).Wait()
216222
require.Len(t, resources, 1)
217223
require.Len(t, resources[0].Agents, 1)
218224
require.Len(t, resources[0].Agents[0].Subsystems, 2)
@@ -221,3 +227,29 @@ func TestWorkspaceAgent(t *testing.T) {
221227
require.Equal(t, codersdk.AgentSubsystemExectrace, resources[0].Agents[0].Subsystems[1])
222228
})
223229
}
230+
231+
func matchAgentWithVersion(rs []codersdk.WorkspaceResource) bool {
232+
if len(rs) < 1 {
233+
return false
234+
}
235+
if len(rs[0].Agents) < 1 {
236+
return false
237+
}
238+
if rs[0].Agents[0].Version == "" {
239+
return false
240+
}
241+
return true
242+
}
243+
244+
func matchAgentWithSubsystems(rs []codersdk.WorkspaceResource) bool {
245+
if len(rs) < 1 {
246+
return false
247+
}
248+
if len(rs[0].Agents) < 1 {
249+
return false
250+
}
251+
if len(rs[0].Agents[0].Subsystems) < 1 {
252+
return false
253+
}
254+
return true
255+
}

‎coderd/agentapi/servicebanner_internal_test.go

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -5,12 +5,12 @@ import (
55
"sync/atomic"
66
"testing"
77

8+
"github.com/stretchr/testify/require"
89
"golang.org/x/xerrors"
910

1011
agentproto "github.com/coder/coder/v2/agent/proto"
1112
"github.com/coder/coder/v2/coderd/appearance"
1213
"github.com/coder/coder/v2/codersdk"
13-
"github.com/stretchr/testify/require"
1414
)
1515

1616
func TestGetServiceBanner(t *testing.T) {

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy