Skip to content

feat: use agent v2 API to update app health #11889

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 1 commit into from
Jan 30, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
3 changes: 1 addition & 2 deletions agent/agent.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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
Expand Down
11 changes: 3 additions & 8 deletions agent/agenttest/client.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
Expand Down Expand Up @@ -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) {
Expand Down
91 changes: 61 additions & 30 deletions coderd/workspaceagents_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down Expand Up @@ -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)
}

Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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{
Expand Down Expand Up @@ -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)
Expand Down
22 changes: 11 additions & 11 deletions codersdk/agentsdk/agentsdk.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
15 changes: 15 additions & 0 deletions codersdk/agentsdk/convert.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}