From 07e8ce01ed7c7471b25f63e04217291d98225173 Mon Sep 17 00:00:00 2001 From: Spike Curtis Date: Mon, 29 Jan 2024 16:54:17 +0400 Subject: [PATCH] feat: use agent v2 API to update app health --- agent/agent.go | 3 +- agent/agenttest/client.go | 11 ++-- coderd/workspaceagents_test.go | 91 +++++++++++++++++++++++----------- codersdk/agentsdk/agentsdk.go | 22 ++++---- codersdk/agentsdk/convert.go | 15 ++++++ 5 files changed, 91 insertions(+), 51 deletions(-) diff --git a/agent/agent.go b/agent/agent.go index bcf4e50cc8bc9..0940b3e04f4ba 100644 --- a/agent/agent.go +++ b/agent/agent.go @@ -91,7 +91,6 @@ type Client interface { Listen(ctx context.Context) (drpc.Conn, error) ReportStats(ctx context.Context, log slog.Logger, statsChan <-chan *agentsdk.Stats, setInterval func(time.Duration)) (io.Closer, error) PostLifecycle(ctx context.Context, state agentsdk.PostLifecycleRequest) error - PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error PostMetadata(ctx context.Context, req agentsdk.PostMetadataRequest) error PatchLogs(ctx context.Context, req agentsdk.PatchLogs) error RewriteDERPMap(derpMap *tailcfg.DERPMap) @@ -808,7 +807,7 @@ func (a *agent) run(ctx context.Context) error { appReporterCtx, appReporterCtxCancel := context.WithCancel(ctx) defer appReporterCtxCancel() go NewWorkspaceAppHealthReporter( - a.logger, manifest.Apps, a.client.PostAppHealth)(appReporterCtx) + a.logger, manifest.Apps, agentsdk.AppHealthPoster(aAPI))(appReporterCtx) a.closeMutex.Lock() network := a.network diff --git a/agent/agenttest/client.go b/agent/agenttest/client.go index 3f721a3135a55..231318c30d65a 100644 --- a/agent/agenttest/client.go +++ b/agent/agenttest/client.go @@ -167,11 +167,6 @@ func (c *Client) PostLifecycle(ctx context.Context, req agentsdk.PostLifecycleRe return nil } -func (c *Client) PostAppHealth(ctx context.Context, req agentsdk.PostAppHealthsRequest) error { - c.logger.Debug(ctx, "post app health", slog.F("req", req)) - return nil -} - func (c *Client) GetStartup() <-chan *agentproto.Startup { return c.fakeAgentAPI.startupCh } @@ -279,9 +274,9 @@ func (*FakeAgentAPI) UpdateLifecycle(context.Context, *agentproto.UpdateLifecycl panic("implement me") } -func (*FakeAgentAPI) BatchUpdateAppHealths(context.Context, *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) { - // TODO implement me - panic("implement me") +func (f *FakeAgentAPI) BatchUpdateAppHealths(ctx context.Context, req *agentproto.BatchUpdateAppHealthRequest) (*agentproto.BatchUpdateAppHealthResponse, error) { + f.logger.Debug(ctx, "batch update app health", slog.F("req", req)) + return &agentproto.BatchUpdateAppHealthResponse{}, nil } func (f *FakeAgentAPI) UpdateStartup(_ context.Context, req *agentproto.UpdateStartupRequest) (*agentproto.Startup, error) { diff --git a/coderd/workspaceagents_test.go b/coderd/workspaceagents_test.go index bf303dd0bf703..cd25d1148f2bf 100644 --- a/coderd/workspaceagents_test.go +++ b/coderd/workspaceagents_test.go @@ -500,7 +500,14 @@ func TestWorkspaceAgentTailnetDirectDisabled(t *testing.T) { // Verify that the manifest has DisableDirectConnections set to true. agentClient := agentsdk.New(client.URL) agentClient.SetSessionToken(r.AgentToken) - manifest := requireGetManifest(ctx, t, agentClient) + rpc, err := agentClient.Listen(ctx) + require.NoError(t, err) + defer func() { + cErr := rpc.Close() + require.NoError(t, cErr) + }() + aAPI := agentproto.NewDRPCAgentClient(rpc) + manifest := requireGetManifest(ctx, t, aAPI) require.True(t, manifest.DisableDirectConnections) _ = agenttest.New(t, client.URL, r.AgentToken) @@ -823,46 +830,63 @@ func TestWorkspaceAgentAppHealth(t *testing.T) { agentClient := agentsdk.New(client.URL) agentClient.SetSessionToken(r.AgentToken) + conn, err := agentClient.Listen(ctx) + require.NoError(t, err) + defer func() { + cErr := conn.Close() + require.NoError(t, cErr) + }() + aAPI := agentproto.NewDRPCAgentClient(conn) - manifest := requireGetManifest(ctx, t, agentClient) + manifest := requireGetManifest(ctx, t, aAPI) require.EqualValues(t, codersdk.WorkspaceAppHealthDisabled, manifest.Apps[0].Health) require.EqualValues(t, codersdk.WorkspaceAppHealthInitializing, manifest.Apps[1].Health) - err := agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{}) - require.Error(t, err) // empty - err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{}) - require.Error(t, err) + _, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{}) + require.NoError(t, err) // healthcheck disabled - err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{ - Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{ - manifest.Apps[0].ID: codersdk.WorkspaceAppHealthInitializing, + _, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{ + Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{ + { + Id: manifest.Apps[0].ID[:], + Health: agentproto.AppHealth_INITIALIZING, + }, }, }) require.Error(t, err) // invalid value - err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{ - Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{ - manifest.Apps[1].ID: codersdk.WorkspaceAppHealth("bad-value"), + _, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{ + Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{ + { + Id: manifest.Apps[1].ID[:], + Health: 99, + }, }, }) require.Error(t, err) // update to healthy - err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{ - Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{ - manifest.Apps[1].ID: codersdk.WorkspaceAppHealthHealthy, + _, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{ + Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{ + { + Id: manifest.Apps[1].ID[:], + Health: agentproto.AppHealth_HEALTHY, + }, }, }) require.NoError(t, err) - manifest = requireGetManifest(ctx, t, agentClient) + manifest = requireGetManifest(ctx, t, aAPI) require.EqualValues(t, codersdk.WorkspaceAppHealthHealthy, manifest.Apps[1].Health) // update to unhealthy - err = agentClient.PostAppHealth(ctx, agentsdk.PostAppHealthsRequest{ - Healths: map[uuid.UUID]codersdk.WorkspaceAppHealth{ - manifest.Apps[1].ID: codersdk.WorkspaceAppHealthUnhealthy, + _, err = aAPI.BatchUpdateAppHealths(ctx, &agentproto.BatchUpdateAppHealthRequest{ + Updates: []*agentproto.BatchUpdateAppHealthRequest_HealthUpdate{ + { + Id: manifest.Apps[1].ID[:], + Health: agentproto.AppHealth_UNHEALTHY, + }, }, }) require.NoError(t, err) - manifest = requireGetManifest(ctx, t, agentClient) + manifest = requireGetManifest(ctx, t, aAPI) require.EqualValues(t, codersdk.WorkspaceAppHealthUnhealthy, manifest.Apps[1].Health) } @@ -1105,8 +1129,15 @@ func TestWorkspaceAgent_Metadata(t *testing.T) { agentClient.SetSessionToken(r.AgentToken) ctx := testutil.Context(t, testutil.WaitMedium) + conn, err := agentClient.Listen(ctx) + require.NoError(t, err) + defer func() { + cErr := conn.Close() + require.NoError(t, cErr) + }() + aAPI := agentproto.NewDRPCAgentClient(conn) - manifest := requireGetManifest(ctx, t, agentClient) + manifest := requireGetManifest(ctx, t, aAPI) // Verify manifest API response. require.Equal(t, workspace.ID, manifest.WorkspaceID) @@ -1276,8 +1307,15 @@ func TestWorkspaceAgent_Metadata_CatchMemoryLeak(t *testing.T) { agentClient.SetSessionToken(r.AgentToken) ctx, cancel := context.WithCancel(testutil.Context(t, testutil.WaitSuperLong)) + conn, err := agentClient.Listen(ctx) + require.NoError(t, err) + defer func() { + cErr := conn.Close() + require.NoError(t, cErr) + }() + aAPI := agentproto.NewDRPCAgentClient(conn) - manifest := requireGetManifest(ctx, t, agentClient) + manifest := requireGetManifest(ctx, t, aAPI) post := func(ctx context.Context, key, value string) error { return agentClient.PostMetadata(ctx, agentsdk.PostMetadataRequest{ @@ -1622,14 +1660,7 @@ func TestWorkspaceAgentExternalAuthListen(t *testing.T) { }) } -func requireGetManifest(ctx context.Context, t testing.TB, client agent.Client) agentsdk.Manifest { - conn, err := client.Listen(ctx) - require.NoError(t, err) - defer func() { - cErr := conn.Close() - require.NoError(t, cErr) - }() - aAPI := agentproto.NewDRPCAgentClient(conn) +func requireGetManifest(ctx context.Context, t testing.TB, aAPI agentproto.DRPCAgentClient) agentsdk.Manifest { mp, err := aAPI.GetManifest(ctx, &agentproto.GetManifestRequest{}) require.NoError(t, err) manifest, err := agentsdk.ManifestFromProto(mp) diff --git a/codersdk/agentsdk/agentsdk.go b/codersdk/agentsdk/agentsdk.go index cc19dd88e3ca9..78429bcfcc083 100644 --- a/codersdk/agentsdk/agentsdk.go +++ b/codersdk/agentsdk/agentsdk.go @@ -231,18 +231,18 @@ type PostAppHealthsRequest struct { Healths map[uuid.UUID]codersdk.WorkspaceAppHealth } -// PostAppHealth updates the workspace agent app health status. -func (c *Client) PostAppHealth(ctx context.Context, req PostAppHealthsRequest) error { - res, err := c.SDK.Request(ctx, http.MethodPost, "/api/v2/workspaceagents/me/app-health", req) - if err != nil { - return err - } - defer res.Body.Close() - if res.StatusCode != http.StatusOK { - return codersdk.ReadBodyAsError(res) +func AppHealthPoster(aAPI proto.DRPCAgentClient) func(ctx context.Context, req PostAppHealthsRequest) error { + return func(ctx context.Context, req PostAppHealthsRequest) error { + pReq, err := ProtoFromAppHealthsRequest(req) + if err != nil { + return xerrors.Errorf("convert AppHealthsRequest: %w", err) + } + _, err = aAPI.BatchUpdateAppHealths(ctx, pReq) + if err != nil { + return xerrors.Errorf("batch update app healths: %w", err) + } + return nil } - - return nil } // AuthenticateResponse is returned when an instance ID diff --git a/codersdk/agentsdk/convert.go b/codersdk/agentsdk/convert.go index 1838e52b013d9..5b556fd688984 100644 --- a/codersdk/agentsdk/convert.go +++ b/codersdk/agentsdk/convert.go @@ -278,3 +278,18 @@ func ProtoFromSubsystems(ss []codersdk.AgentSubsystem) ([]proto.Startup_Subsyste } return ret, nil } + +func ProtoFromAppHealthsRequest(req PostAppHealthsRequest) (*proto.BatchUpdateAppHealthRequest, error) { + pReq := &proto.BatchUpdateAppHealthRequest{} + for id, h := range req.Healths { + hp, ok := proto.AppHealth_value[strings.ToUpper(string(h))] + if !ok { + return nil, xerrors.Errorf("unknown app health: %s", h) + } + pReq.Updates = append(pReq.Updates, &proto.BatchUpdateAppHealthRequest_HealthUpdate{ + Id: id[:], + Health: proto.AppHealth(hp), + }) + } + return pReq, nil +}