Movatterモバイル変換

[0]ホーム
URL:

Skip to content

Navigation Menu

Sign in
Appearance settings

Search code, repositories, users, issues, pull requests...

Provide feedback

We read every piece of feedback, and take your input very seriously.

Saved searches

Use saved searches to filter your results more quickly

 Sign up
Appearance settings

Commit04cdbc1

Browse files
committed
feat: use Agent v2 API for Service Banner
1 parent174664c commit04cdbc1

File tree

7 files changed

+264
-127
lines changed

7 files changed

+264
-127
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
typeAgentinterface {
@@ -269,7 +269,6 @@ func (a *agent) init(ctx context.Context) {
269269
func (a*agent)runLoop(ctx context.Context) {
270270
goa.reportLifecycleLoop(ctx)
271271
goa.reportMetadataLoop(ctx)
272-
goa.fetchServiceBannerLoop(ctx)
273272
goa.manageProcessPriorityLoop(ctx)
274273

275274
forretrier:=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
deferticker.Stop()
668667
for {
669668
select {
670669
case<-ctx.Done():
671-
return
670+
returnctx.Err()
672671
case<-ticker.C:
673-
serviceBanner,err:=a.client.GetServiceBanner(ctx)
672+
sbp,err:=aAPI.GetServiceBanner(ctx,&proto.GetServiceBannerRequest{})
674673
iferr!=nil {
675674
ifctx.Err()!=nil {
676-
return
675+
returnctx.Err()
677676
}
678677
a.logger.Error(ctx,"failed to update service banner",slog.Error(err))
679-
continue
678+
returnerr
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+
iferr!=nil {
699+
returnerr
700+
}
701+
deferfunc() {
702+
cErr:=conn.Close()
703+
ifcErr!=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
iferr!=nil {
698711
returnxerrors.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-
iferr!=nil {
827-
returnerr
828-
}
829-
deferfunc() {
830-
cErr:=conn.Close()
831-
ifcErr!=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
returnnil
853855
})
854856

857+
eg.Go(func()error {
858+
a.logger.Debug(egCtx,"running fetch server banner loop")
859+
err:=a.fetchServiceBannerLoop(egCtx,aAPI)
860+
iferr!=nil {
861+
returnxerrors.Errorf("fetch server banner loop: %w",err)
862+
}
863+
returnnil
864+
})
865+
855866
returneg.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(errerror) {
5357
ifxerrors.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
typeClientstruct {
72-
ttesting.TB
73-
loggerslog.Logger
74-
agentIDuuid.UUID
75-
manifestagentsdk.Manifest
76-
metadatamap[string]agentsdk.Metadata
77-
statsChanchan*agentsdk.Stats
78-
coordinatortailnet.Coordinator
79-
server*drpcserver.Server
80-
LastWorkspaceAgentfunc()
81-
PatchWorkspaceLogsfunc()error
82-
GetServiceBannerFuncfunc()(codersdk.ServiceBannerConfig,error)
77+
t testing.TB
78+
logger slog.Logger
79+
agentID uuid.UUID
80+
manifest agentsdk.Manifest
81+
metadatamap[string]agentsdk.Metadata
82+
statsChanchan*agentsdk.Stats
83+
coordinator tailnet.Coordinator
84+
server*drpcserver.Server
85+
fakeAgentAPI*FakeAgentAPI
86+
LastWorkspaceAgentfunc()
87+
PatchWorkspaceLogsfunc()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(ffunc() (codersdk.ServiceBannerConfig,error)) {
224-
c.mu.Lock()
225-
deferc.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-
deferc.mu.Unlock()
233-
c.logger.Debug(ctx,"get service banner")
234-
ifc.GetServiceBannerFunc!=nil {
235-
returnc.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 (ccloseFunc)Close()error {
255247
returnc()
256248
}
249+
250+
typeFakeAgentAPIstruct {
251+
sync.Mutex
252+
t testing.TB
253+
logger slog.Logger
254+
255+
getServiceBannerFuncfunc() (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(fnfunc() (codersdk.ServiceBannerConfig,error)) {
264+
f.Lock()
265+
deferf.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+
deferf.Unlock()
273+
iff.getServiceBannerFunc==nil {
274+
return&agentproto.ServiceBanner{},nil
275+
}
276+
sb,err:=f.getServiceBannerFunc()
277+
iferr!=nil {
278+
returnnil,err
279+
}
280+
returnagentproto.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+
funcNewFakeAgentAPI(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+
funcmatchAgentWithVersion(rs []codersdk.WorkspaceResource)bool {
232+
iflen(rs)<1 {
233+
returnfalse
234+
}
235+
iflen(rs[0].Agents)<1 {
236+
returnfalse
237+
}
238+
ifrs[0].Agents[0].Version=="" {
239+
returnfalse
240+
}
241+
returntrue
242+
}
243+
244+
funcmatchAgentWithSubsystems(rs []codersdk.WorkspaceResource)bool {
245+
iflen(rs)<1 {
246+
returnfalse
247+
}
248+
iflen(rs[0].Agents)<1 {
249+
returnfalse
250+
}
251+
iflen(rs[0].Agents[0].Subsystems)<1 {
252+
returnfalse
253+
}
254+
returntrue
255+
}

0 commit comments

Comments
 (0)
[8]ページ先頭
©2009-2025 Movatter.jp