Skip to content

Commit 87bc524

Browse files
jiandongPKU Git
jiandong
authored and
PKU Git
committed
Merge pull request 'feat(api): manage command/event/health lifecycle' (#2) from feat/command-retention into main
2 parents d0d0047 + dd52efe commit 87bc524

11 files changed

Lines changed: 129 additions & 30 deletions

‎pkg/api/agent_service.go‎

Lines changed: 7 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@ package api
33
import (
44
"context"
55
"fmt"
6+
"log/slog"
67
"sort"
78
"time"
89

@@ -264,9 +265,7 @@ func (s *AgentService) RecordCommandEvent(ctx context.Context, event controlplan
264265
if s.registry != nil {
265266
s.registry.SaveCommandEvent(event)
266267
}
267-
if s.store != nil {
268-
_ = s.store.SaveAgentCommandEvent(ctx, event)
269-
}
268+
slog.Info("agent command event", "agent_id", event.AgentID, "command_id", event.CommandID, "type", event.Type, "status", event.Status)
270269
}
271270

272271
func (s *AgentService) updateCommandFromEvent(ctx context.Context, event controlplane.AgentCommandEvent) {
@@ -301,7 +300,11 @@ func (s *AgentService) updateCommandFromEvent(ctx context.Context, event control
301300
_ = newGuestFileOperationService(s.store).ReconcileCommandFailure(ctx, cmd.FilePut.ID, event.Status)
302301
}
303302
}
304-
_ = s.store.SaveAgentCommand(ctx, cmd)
303+
if cmd.IsTerminal() {
304+
_ = s.store.DeleteAgentCommand(ctx, cmd.AgentID, cmd.ID)
305+
} else {
306+
_ = s.store.SaveAgentCommand(ctx, cmd)
307+
}
305308
}
306309

307310
func (s *AgentService) updateConsoleSessionFromCommandEvent(cmd controlplane.AgentCommand, event controlplane.AgentCommandEvent) {

‎pkg/api/agent_service_test.go‎

Lines changed: 14 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -52,7 +52,7 @@ func TestAgentServicePublishAndAcquireCommandLease(t *testing.T) {
5252
}
5353

5454
func TestAgentServiceRecordCommandEventUpdatesCommand(t *testing.T) {
55-
svc, store, registry, _ := newTestAgentService(t)
55+
svc, _, registry, _ := newTestAgentService(t)
5656
ctx := context.Background()
5757
cmd, err := svc.PublishCommand(ctx, "host-a", controlplane.AgentCommand{Type: "node_operation"})
5858
require.NoError(t, err)
@@ -61,19 +61,25 @@ func TestAgentServiceRecordCommandEventUpdatesCommand(t *testing.T) {
6161
svc.RecordCommandEvent(ctx, controlplane.AgentCommandEvent{
6262
AgentID: "host-a",
6363
CommandID: cmd.ID,
64-
Status: controlplane.AgentCommandStatusCompleted,
64+
Status: "ack",
6565
CreatedAt: now,
6666
})
6767

6868
got, err := svc.FindCommand(ctx, "host-a", cmd.ID)
6969
require.NoError(t, err)
70-
require.Equal(t, controlplane.AgentCommandStatusCompleted, got.Status)
71-
require.Equal(t, now, got.EndedAt)
70+
require.Equal(t, controlplane.AgentCommandStatusAcknowledged, got.Status)
7271

73-
events, err := store.ListAgentCommandEvents(ctx, "host-a")
74-
require.NoError(t, err)
75-
require.Len(t, events, 1)
76-
require.Len(t, registry.ListCommandEvents("host-a"), 1)
72+
// A terminal event removes the command rather than keeping a completed row.
73+
svc.RecordCommandEvent(ctx, controlplane.AgentCommandEvent{
74+
AgentID: "host-a",
75+
CommandID: cmd.ID,
76+
Status: controlplane.AgentCommandStatusCompleted,
77+
CreatedAt: now,
78+
})
79+
_, err = svc.FindCommand(ctx, "host-a", cmd.ID)
80+
require.Error(t, err)
81+
82+
require.Len(t, registry.ListCommandEvents("host-a"), 2)
7783
}
7884

7985
func TestAgentServiceReconcilesStalePendingCommandFromGuestOperation(t *testing.T) {

‎pkg/api/api_store.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -207,6 +207,14 @@ func (s *localAPIStore) SaveHealth(_ context.Context, topology string, snap Heal
207207
return os.WriteFile(path, raw, 0o644)
208208
}
209209

210+
func (s *localAPIStore) DeleteHealth(_ context.Context, topology string) error {
211+
path := filepath.Join(s.runsDir, topology, "health.json")
212+
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
213+
return err
214+
}
215+
return nil
216+
}
217+
210218
func (s *localAPIStore) LoadHealth(_ context.Context, topology string) (*HealthSnapshot, error) {
211219
raw, err := os.ReadFile(filepath.Join(s.runsDir, topology, "health.json"))
212220
if err != nil {
@@ -534,6 +542,18 @@ ON CONFLICT (topology) DO UPDATE SET data=EXCLUDED.data, updated_at=now()`,
534542
return nil
535543
}
536544

545+
func (s *postgresAPIStore) DeleteHealth(ctx context.Context, topology string) error {
546+
conn, err := s.connect(ctx)
547+
if err != nil {
548+
return err
549+
}
550+
defer conn.Release()
551+
if _, err := conn.Exec(ctx, `DELETE FROM sysbox_health WHERE topology=$1`, topology); err != nil {
552+
return fmt.Errorf("postgres delete health: %w", err)
553+
}
554+
return nil
555+
}
556+
537557
func (s *postgresAPIStore) LoadHealth(ctx context.Context, topology string) (*HealthSnapshot, error) {
538558
conn, err := s.connect(ctx)
539559
if err != nil {

‎pkg/api/api_store_agents.go‎

Lines changed: 20 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -74,6 +74,14 @@ func (s *localAPIStore) SaveAgentCommand(_ context.Context, cmd controlplane.Age
7474
return writeLocalObject(path, cmd)
7575
}
7676

77+
func (s *localAPIStore) DeleteAgentCommand(_ context.Context, agentID, commandID string) error {
78+
path := filepath.Join(s.runsDir, "_agents", agentID, "commands", commandID+".json")
79+
if err := os.Remove(path); err != nil && !os.IsNotExist(err) {
80+
return err
81+
}
82+
return nil
83+
}
84+
7785
func (s *localAPIStore) ListAgentCommands(_ context.Context, agentID string) ([]controlplane.AgentCommand, error) {
7886
var commands []controlplane.AgentCommand
7987
var err error
@@ -342,6 +350,18 @@ WHERE
342350
return nil
343351
}
344352

353+
func (s *postgresAPIStore) DeleteAgentCommand(ctx context.Context, agentID, commandID string) error {
354+
conn, err := s.connect(ctx)
355+
if err != nil {
356+
return err
357+
}
358+
defer conn.Release()
359+
if _, err := conn.Exec(ctx, `DELETE FROM sysbox_agent_commands WHERE id=$1 AND agent_id=$2`, commandID, agentID); err != nil {
360+
return fmt.Errorf("postgres delete agent command: %w", err)
361+
}
362+
return nil
363+
}
364+
345365
func (s *postgresAPIStore) AcquireAgentCommandLease(ctx context.Context, agentID, commandID, owner string, ttl time.Duration) (*controlplane.AgentCommand, bool, error) {
346366
conn, err := s.connect(ctx)
347367
if err != nil {

‎pkg/api/api_store_interfaces.go‎

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -36,6 +36,7 @@ type checkpointStore interface {
3636
type healthStore interface {
3737
SaveHealth(ctx context.Context, topology string, snap HealthSnapshot) error
3838
LoadHealth(ctx context.Context, topology string) (*HealthSnapshot, error)
39+
DeleteHealth(ctx context.Context, topology string) error
3940
}
4041

4142
type revisionStore interface {
@@ -77,6 +78,7 @@ type agentCommandStore interface {
7778
SaveAgentCommandEvent(ctx context.Context, event controlplane.AgentCommandEvent) error
7879
ListAgentCommandEvents(ctx context.Context, agentID string) ([]controlplane.AgentCommandEvent, error)
7980
SaveAgentCommand(ctx context.Context, cmd controlplane.AgentCommand) error
81+
DeleteAgentCommand(ctx context.Context, agentID, commandID string) error
8082
ListAgentCommands(ctx context.Context, agentID string) ([]controlplane.AgentCommand, error)
8183
AcquireAgentCommandLease(ctx context.Context, agentID, commandID, owner string, ttl time.Duration) (*controlplane.AgentCommand, bool, error)
8284
}

‎pkg/api/api_store_sqlite.go‎

Lines changed: 18 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -527,6 +527,15 @@ func (s *sqliteAPIStore) SaveHealth(ctx context.Context, topology string, snap H
527527
return err
528528
}
529529

530+
func (s *sqliteAPIStore) DeleteHealth(ctx context.Context, topology string) error {
531+
db, err := s.open()
532+
if err != nil {
533+
return err
534+
}
535+
_, err = db.ExecContext(ctx, `DELETE FROM sysbox_health WHERE topology=?`, topology)
536+
return err
537+
}
538+
530539
func (s *sqliteAPIStore) LoadHealth(ctx context.Context, topology string) (*HealthSnapshot, error) {
531540
db, err := s.open()
532541
if err != nil {
@@ -1015,6 +1024,15 @@ func (s *sqliteAPIStore) SaveAgentCommand(ctx context.Context, cmd controlplane.
10151024
return err
10161025
}
10171026

1027+
func (s *sqliteAPIStore) DeleteAgentCommand(ctx context.Context, agentID, commandID string) error {
1028+
db, err := s.open()
1029+
if err != nil {
1030+
return err
1031+
}
1032+
_, err = db.ExecContext(ctx, `DELETE FROM sysbox_agent_commands WHERE id=? AND agent_id=?`, commandID, agentID)
1033+
return err
1034+
}
1035+
10181036
func (s *sqliteAPIStore) ListAgentCommands(ctx context.Context, agentID string) ([]controlplane.AgentCommand, error) {
10191037
db, err := s.open()
10201038
if err != nil {

‎pkg/api/guest_file_test.go‎

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -200,16 +200,24 @@ func TestCancelIntentBlocksConcurrentStartAndCompletion(t *testing.T) {
200200

201201
func TestGuestFileFailedCommandEventReconcilesDurableOperation(t *testing.T) {
202202
s := guestExecutionTestServer(t)
203-
op := controlplane.GuestFileOperation{ID: "op", Version: 1, AgentID: "host-a", CommandID: "cmd", Status: controlplane.GuestExecutionRunning}
204-
require.NoError(t, s.apiStore.SaveGuestFileOperation(context.Background(), op))
205-
_, err := s.agentService().PublishCommand(context.Background(), "host-a", controlplane.AgentCommand{ID: "cmd", Type: "guest_file_put", FilePut: &controlplane.GuestFilePut{ID: "op"}})
203+
204+
// A completed command leaves the durable operation running.
205+
done := controlplane.GuestFileOperation{ID: "op-done", Version: 1, AgentID: "host-a", CommandID: "cmd-done", Status: controlplane.GuestExecutionRunning}
206+
require.NoError(t, s.apiStore.SaveGuestFileOperation(context.Background(), done))
207+
_, err := s.agentService().PublishCommand(context.Background(), "host-a", controlplane.AgentCommand{ID: "cmd-done", Type: "guest_file_put", FilePut: &controlplane.GuestFilePut{ID: "op-done"}})
206208
require.NoError(t, err)
207-
s.agentService().RecordCommandEvent(context.Background(), controlplane.AgentCommandEvent{CommandID: "cmd", AgentID: "host-a", Status: controlplane.AgentCommandStatusCompleted})
208-
got, err := s.apiStore.GetGuestFileOperation(context.Background(), "op")
209+
s.agentService().RecordCommandEvent(context.Background(), controlplane.AgentCommandEvent{CommandID: "cmd-done", AgentID: "host-a", Status: controlplane.AgentCommandStatusCompleted})
210+
got, err := s.apiStore.GetGuestFileOperation(context.Background(), "op-done")
209211
require.NoError(t, err)
210212
require.Equal(t, controlplane.GuestExecutionRunning, got.Status)
211-
s.agentService().RecordCommandEvent(context.Background(), controlplane.AgentCommandEvent{CommandID: "cmd", AgentID: "host-a", Status: controlplane.AgentCommandStatusFailed})
212-
got, err = s.apiStore.GetGuestFileOperation(context.Background(), "op")
213+
214+
// A failed command reconciles the durable operation to failed.
215+
failed := controlplane.GuestFileOperation{ID: "op-failed", Version: 1, AgentID: "host-a", CommandID: "cmd-failed", Status: controlplane.GuestExecutionRunning}
216+
require.NoError(t, s.apiStore.SaveGuestFileOperation(context.Background(), failed))
217+
_, err = s.agentService().PublishCommand(context.Background(), "host-a", controlplane.AgentCommand{ID: "cmd-failed", Type: "guest_file_put", FilePut: &controlplane.GuestFilePut{ID: "op-failed"}})
218+
require.NoError(t, err)
219+
s.agentService().RecordCommandEvent(context.Background(), controlplane.AgentCommandEvent{CommandID: "cmd-failed", AgentID: "host-a", Status: controlplane.AgentCommandStatusFailed})
220+
got, err = s.apiStore.GetGuestFileOperation(context.Background(), "op-failed")
213221
require.NoError(t, err)
214222
require.Equal(t, controlplane.GuestExecutionFailed, got.Status)
215223
}

‎pkg/api/server.go‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ func NewServerWithConfig(cfg config.ServiceConfig) *Server {
7575
mux: http.NewServeMux(),
7676
}
7777
s.agentSvc = newAgentService(s)
78-
s.workspaces = newWorkspaceService(s.runsDir, s.workspacesDir, s.stateBackend, s.stateManager)
78+
s.workspaces = newWorkspaceService(s.runsDir, s.workspacesDir, s.stateBackend, s.stateManager, s.apiStore)
7979
s.planService = newPlanService(s)
8080
s.scheduler = newSchedulerService(s)
8181
s.nodeService = newNodeOperationService(s.workspaceService(), s.scheduling(), s.nodeOps, s.agentService().PublishCommand)
@@ -134,7 +134,7 @@ func (s *Server) scheduling() *SchedulerService {
134134

135135
func (s *Server) workspaceService() *WorkspaceService {
136136
if s.workspaces == nil {
137-
s.workspaces = newWorkspaceService(s.runsDir, s.workspacesDir, s.stateBackend, s.stateManager)
137+
s.workspaces = newWorkspaceService(s.runsDir, s.workspacesDir, s.stateBackend, s.stateManager, s.apiStore)
138138
}
139139
return s.workspaces
140140
}

‎pkg/api/workspace_service.go‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -17,6 +17,7 @@ type WorkspaceService struct {
1717
workspacesDir string
1818
stateBackend string
1919
stateManager func(topology string) (*state.Manager, error)
20+
health healthStore
2021
}
2122

2223
type WorkspaceInfo struct {
@@ -30,12 +31,13 @@ type WorkspaceInfo struct {
3031
Backend string `json:"backend,omitempty"`
3132
}
3233

33-
func newWorkspaceService(runsDir, workspacesDir, stateBackend string, stateManager func(string) (*state.Manager, error)) *WorkspaceService {
34+
func newWorkspaceService(runsDir, workspacesDir, stateBackend string, stateManager func(string) (*state.Manager, error), health healthStore) *WorkspaceService {
3435
return &WorkspaceService{
3536
runsDir: runsDir,
3637
workspacesDir: workspacesDir,
3738
stateBackend: stateBackend,
3839
stateManager: stateManager,
40+
health: health,
3941
}
4042
}
4143

@@ -198,6 +200,9 @@ func (s *WorkspaceService) Delete(ctx context.Context, topology string, force bo
198200
if err := os.RemoveAll(filepath.Dir(s.HCLFile(topology))); err != nil {
199201
return fmt.Errorf("remove workspace: %w", err)
200202
}
203+
if s.health != nil {
204+
_ = s.health.DeleteHealth(ctx, topology)
205+
}
201206
return nil
202207
}
203208

‎pkg/provider/docker/nic.go‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -64,6 +64,15 @@ func (s *Substrate) Attach(ctx context.Context, h substrate.NodeHandle, req driv
6464
if err != nil {
6565
return driver.AttachmentResult{}, driver.Wrap(driver.ErrorUnavailable, "docker", "attach network", err)
6666
}
67+
if target.NAT {
68+
// Docker names the interface on managed networks; resolve it so NAT
69+
// policy can bind rules to the actual guest device.
70+
resolved, err := s.resolveAttachmentDevice(ctx, h, req, driver.AttachmentResult{})
71+
if err != nil {
72+
return driver.AttachmentResult{}, driver.Wrap(driver.ErrorUnavailable, "docker", "resolve NAT attachment device", err)
73+
}
74+
guest = resolved
75+
}
6776
if target.NAT {
6877
if hs, ok := h.Provider.(*HandleState); ok && hs.RemoveDefaultBridge {
6978
if err := s.cli.NetworkDisconnect(ctx, "bridge", h.ID, true); err != nil {

0 commit comments

Comments
 (0)