Skip to content

Commit 0fc1772

Browse files
spikecurtisspikecurtis
authored
feat: use agent v2 API to update app health (#11889)
Use the Agent v2 API to update App Health
1 parent 2599850 commit 0fc1772

File tree

5 files changed

+91
-51
lines changed

5 files changed

+91
-51
lines changed

‎agent/agent.go

+1-2
Original file line numberDiff line numberDiff line change
@@ -91,7 +91,6 @@ type Client interface {
9191
Listen(ctx context.Context) (drpc.Conn, error)
9292
ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error)
9393
PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error
94-
PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error
9594
PostMetadata(ctx context.Context, req agentsdk.PostMetadataRequest) error
9695
PatchLogs(ctx context.Context, req agentsdk.PatchLogs) error
9796
RewriteDERPMap(derpMap *tailcfg.DERPMap)
@@ -808,7 +807,7 @@ func (a *agent) run(ctx context.Context) error {
808807
appReporterCtx, appReporterCtxCancel := context.WithCancel(ctx)
809808
defer appReporterCtxCancel()
810809
go NewWorkspaceAppHealthReporter(
811-
a.logger, manifest.Apps, a.client.PostAppHealth)(appReporterCtx)
810+
a.logger, manifest.Apps, agentsdk.AppHealthPoster(aAPI))(appReporterCtx)
812811

813812
a.closeMutex.Lock()
814813
network := a.network

‎agent/agenttest/client.go

+3-8
Original file line numberDiff line numberDiff line change
@@ -167,11 +167,6 @@ func (c *Client) PostLifecycle(ctx context.Context, req agentsdk.PostLifecycleRe
167167
return nil
168168
}
169169

170-
func (c *Client) PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error {
171-
c.logger.Debug(ctx, "post app health", slog.F("req", req))
172-
return nil
173-
}
174-
175170
func (c *Client) GetStartup() <-chan *agentproto.Startup {
176171
return c.fakeAgentAPI.startupCh
177172
}
@@ -279,9 +274,9 @@ func (*FakeAgentAPI) UpdateLifecycle(context.Context, *agentproto.UpdateLifecycl
279274
panic("implement me")
280275
}
281276

282-
func (*FakeAgentAPI) BatchUpdateAppHealths(context.Context, *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
283-
// TODO implement me
284-
panic("implement me")
277+
func (f *FakeAgentAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) {
278+
f.logger.Debug(ctx, "batch update app health", slog.F("req", req))
279+
return &agentproto.BatchUpdateAppHealthResponse{}, nil
285280
}
286281

287282
func (f *FakeAgentAPI) UpdateStartup(_ context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) {

‎coderd/workspaceagents_test.go

+61-30
Original file line numberDiff line numberDiff line change
@@ -500,7 +500,14 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) {
500500
// Verify that the manifest has DisableDirectConnections set to true.
501501
agentClient := agentsdk.New(client.URL)
502502
agentClient.SetSessionToken(r.AgentToken)
503-
manifest := requireGetManifest(ctx, t, agentClient)
503+
rpc, err := agentClient.Listen(ctx)
504+
require.NoError(t, err)
505+
defer func() {
506+
cErr := rpc.Close()
507+
require.NoError(t, cErr)
508+
}()
509+
aAPI := agentproto.NewDRPCAgentClient(rpc)
510+
manifest := requireGetManifest(ctx, t, aAPI)
504511
require.True(t, manifest.DisableDirectConnections)
505512

506513
_ = agenttest.New(t, client.URL, r.AgentToken)
@@ -823,46 +830,63 @@ func TestWorkspaceAgentAppHealth(t *testing.T) {
823830

824831
agentClient := agentsdk.New(client.URL)
825832
agentClient.SetSessionToken(r.AgentToken)
833+
conn, err := agentClient.Listen(ctx)
834+
require.NoError(t, err)
835+
defer func() {
836+
cErr := conn.Close()
837+
require.NoError(t, cErr)
838+
}()
839+
aAPI := agentproto.NewDRPCAgentClient(conn)
826840

827-
manifest := requireGetManifest(ctx, t, agentClient)
841+
manifest := requireGetManifest(ctx, t, aAPI)
828842
require.EqualValues(t, codersdk.WorkspaceAppHealthDisabled, manifest.Apps[0].Health)
829843
require.EqualValues(t, codersdk.WorkspaceAppHealthInitializing, manifest.Apps[1].Health)
830-
err := agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{})
831-
require.Error(t, err)
832844
// empty
833-
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{})
834-
require.Error(t, err)
845+
_, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{})
846+
require.NoError(t, err)
835847
// healthcheck disabled
836-
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{
837-
Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{
838-
manifest.Apps[0].ID: codersdk.WorkspaceAppHealthInitializing,
848+
_, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{
849+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
850+
{
851+
Id: manifest.Apps[0].ID[:],
852+
Health: agentproto.AppHealth_INITIALIZING,
853+
},
839854
},
840855
})
841856
require.Error(t, err)
842857
// invalid value
843-
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{
844-
Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{
845-
manifest.Apps[1].ID: codersdk.WorkspaceAppHealth("bad-value"),
858+
_, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{
859+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
860+
{
861+
Id: manifest.Apps[1].ID[:],
862+
Health: 99,
863+
},
846864
},
847865
})
848866
require.Error(t, err)
849867
// update to healthy
850-
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{
851-
Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{
852-
manifest.Apps[1].ID: codersdk.WorkspaceAppHealthHealthy,
868+
_, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{
869+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
870+
{
871+
Id: manifest.Apps[1].ID[:],
872+
Health: agentproto.AppHealth_HEALTHY,
873+
},
853874
},
854875
})
855876
require.NoError(t, err)
856-
manifest = requireGetManifest(ctx, t, agentClient)
877+
manifest = requireGetManifest(ctx, t, aAPI)
857878
require.EqualValues(t, codersdk.WorkspaceAppHealthHealthy, manifest.Apps[1].Health)
858879
// update to unhealthy
859-
err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{
860-
Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{
861-
manifest.Apps[1].ID: codersdk.WorkspaceAppHealthUnhealthy,
880+
_, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{
881+
Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{
882+
{
883+
Id: manifest.Apps[1].ID[:],
884+
Health: agentproto.AppHealth_UNHEALTHY,
885+
},
862886
},
863887
})
864888
require.NoError(t, err)
865-
manifest = requireGetManifest(ctx, t, agentClient)
889+
manifest = requireGetManifest(ctx, t, aAPI)
866890
require.EqualValues(t, codersdk.WorkspaceAppHealthUnhealthy, manifest.Apps[1].Health)
867891
}
868892

@@ -1105,8 +1129,15 @@ func TestWorkspaceAgent_Metadata(t *testing.T) {
11051129
agentClient.SetSessionToken(r.AgentToken)
11061130

11071131
ctx := testutil.Context(t, testutil.WaitMedium)
1132+
conn, err := agentClient.Listen(ctx)
1133+
require.NoError(t, err)
1134+
defer func() {
1135+
cErr := conn.Close()
1136+
require.NoError(t, cErr)
1137+
}()
1138+
aAPI := agentproto.NewDRPCAgentClient(conn)
11081139

1109-
manifest := requireGetManifest(ctx, t, agentClient)
1140+
manifest := requireGetManifest(ctx, t, aAPI)
11101141

11111142
// Verify manifest API response.
11121143
require.Equal(t, workspace.ID, manifest.WorkspaceID)
@@ -1276,8 +1307,15 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) {
12761307
agentClient.SetSessionToken(r.AgentToken)
12771308

12781309
ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitSuperLong))
1310+
conn, err := agentClient.Listen(ctx)
1311+
require.NoError(t, err)
1312+
defer func() {
1313+
cErr := conn.Close()
1314+
require.NoError(t, cErr)
1315+
}()
1316+
aAPI := agentproto.NewDRPCAgentClient(conn)
12791317

1280-
manifest := requireGetManifest(ctx, t, agentClient)
1318+
manifest := requireGetManifest(ctx, t, aAPI)
12811319

12821320
post := func(ctx context.Context, key, value string) error {
12831321
return agentClient.PostMetadata(ctx, agentsdk.PostMetadataRequest{
@@ -1622,14 +1660,7 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) {
16221660
})
16231661
}
16241662

1625-
func requireGetManifest(ctx context.Context, t testing.TB, client agent.Client) agentsdk.Manifest {
1626-
conn, err := client.Listen(ctx)
1627-
require.NoError(t, err)
1628-
defer func() {
1629-
cErr := conn.Close()
1630-
require.NoError(t, cErr)
1631-
}()
1632-
aAPI := agentproto.NewDRPCAgentClient(conn)
1663+
func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCAgentClient) agentsdk.Manifest {
16331664
mp, err := aAPI.GetManifest(ctx, &agentproto.GetManifestRequest{})
16341665
require.NoError(t, err)
16351666
manifest, err := agentsdk.ManifestFromProto(mp)

‎codersdk/agentsdk/agentsdk.go

+11-11
Original file line numberDiff line numberDiff line change
@@ -231,18 +231,18 @@ type PostAppHealthsRequest struct {
231231
Healths map[uuid.UUID]codersdk.WorkspaceAppHealth
232232
}
233233

234-
// PostAppHealth updates the workspace agent app health status.
235-
func (c *Client) PostAppHealth(ctx context.Context, req PostAppHealthsRequest) error {
236-
res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/app-health", req)
237-
if err != nil {
238-
return err
239-
}
240-
defer res.Body.Close()
241-
if res.StatusCode != http.StatusOK {
242-
return codersdk.ReadBodyAsError(res)
234+
func AppHealthPoster(aAPI proto.DRPCAgentClient) func(ctx context.Context, req PostAppHealthsRequest) error {
235+
return func(ctx context.Context, req PostAppHealthsRequest) error {
236+
pReq, err := ProtoFromAppHealthsRequest(req)
237+
if err != nil {
238+
return xerrors.Errorf("convert AppHealthsRequest: %w", err)
239+
}
240+
_, err = aAPI.BatchUpdateAppHealths(ctx, pReq)
241+
if err != nil {
242+
return xerrors.Errorf("batch update app healths: %w", err)
243+
}
244+
return nil
243245
}
244-
245-
return nil
246246
}
247247

248248
// AuthenticateResponse is returned when an instance ID

‎codersdk/agentsdk/convert.go

+15
Original file line numberDiff line numberDiff line change
@@ -278,3 +278,18 @@ func ProtoFromSubsystems(ss []codersdk.AgentSubsystem) ([]proto.Startup_Subsyste
278278
}
279279
return ret, nil
280280
}
281+
282+
func ProtoFromAppHealthsRequest(req PostAppHealthsRequest) (*proto.BatchUpdateAppHealthRequest, error) {
283+
pReq := &proto.BatchUpdateAppHealthRequest{}
284+
for id, h := range req.Healths {
285+
hp, ok := proto.AppHealth_value[strings.ToUpper(string(h))]
286+
if !ok {
287+
return nil, xerrors.Errorf("unknown app health: %s", h)
288+
}
289+
pReq.Updates = append(pReq.Updates, &proto.BatchUpdateAppHealthRequest_HealthUpdate{
290+
Id: id[:],
291+
Health: proto.AppHealth(hp),
292+
})
293+
}
294+
return pReq, nil
295+
}

0 commit comments

Comments
 (0)