feat: agent reports workloads and handles workload commands

2026-08-07 08:58:46 +01:00
parent cf9d85b3cd
commit 1b351cfca4
3 changed files with 185 additions and 0 deletions
@@ -113,6 +113,19 @@ func (c *Client) ReportPackages(req *pb.ReportPackagesRequest) (bool, error) {
 	return resp.NeedFull, nil
 }

+// ReportWorkloads sends a workload report and returns whether the server wants
+// the full list.
+func (c *Client) ReportWorkloads(req *pb.ReportWorkloadsRequest) (bool, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 60*time.Second)
+	defer cancel()
+
+	resp, err := c.client.ReportWorkloads(ctx, req)
+	if err != nil {
+		return false, err
+	}
+	return resp.NeedFull, nil
+}
+
 func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, privateKey, label string) (string, error) {
 	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
 	defer cancel()
@@ -70,6 +70,8 @@ func Run(ctx context.Context, cfg *config.Config, version string) error {

 	go runInventory(ctx, cfg)

+	go runWorkloads(ctx, cfg)
+
 	go monitors.Run(ctx, cfg)

 	ticker := time.NewTicker(cfg.PollInterval)
@@ -347,6 +349,15 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error {
 		if cmd.OpenProxy != nil {
 			go handleOpenProxy(ctx, cfg, cmd.OpenProxy)
 		}
+		if cmd.RefreshWorkloads != nil {
+			go handleRefreshWorkloads(cfg)
+		}
+		if cmd.ControlWorkload != nil {
+			go handleControlWorkload(send, cfg, cmd.CommandId, cmd.ControlWorkload)
+		}
+		if cmd.WorkloadLogs != nil {
+			go handleWorkloadLogs(send, cfg, cmd.CommandId, cmd.WorkloadLogs)
+		}
 		if cmd.RunStep != nil {
 			go func(rc *pb.RunStepCmd, cid string) {
 				emit := func(seq uint64, data []byte) {
@@ -0,0 +1,161 @@
+package agentsync
+
+import (
+	"context"
+	"log"
+	"runtime"
+	"time"
+
+	"gitea.hostxtra.co.uk/mrhid6/vantage/agent/internal/config"
+	grpcclient "gitea.hostxtra.co.uk/mrhid6/vantage/agent/internal/grpc"
+	"gitea.hostxtra.co.uk/mrhid6/vantage/agent/internal/grpc/pb"
+	"gitea.hostxtra.co.uk/mrhid6/vantage/agent/internal/workloads"
+)
+
+// workloadInterval is the report cadence. Sixty seconds is affordable because
+// an unchanged list costs one small offer message, not the body.
+const workloadInterval = 60 * time.Second
+
+// runWorkloads reports what this host runs, on its own ticker.
+func runWorkloads(ctx context.Context, cfg *config.Config) {
+	if runtime.GOOS != "linux" {
+		return
+	}
+
+	reportWorkloads(cfg)
+
+	ticker := time.NewTicker(workloadInterval)
+	defer ticker.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-ticker.C:
+			reportWorkloads(cfg)
+		}
+	}
+}
+
+// reportWorkloads offers a hash of the current workload set and sends the full
+// list only if the server does not already hold it.
+//
+// This is the ONLY writer of the server_workloads collection. RefreshWorkloadsCmd
+// calls straight into here rather than answering with data of its own.
+func reportWorkloads(cfg *config.Config) {
+	if runtime.GOOS != "linux" {
+		return
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 2*time.Minute)
+	defer cancel()
+
+	res := workloads.Collect(ctx)
+	hash := workloads.Hash(res.Workloads)
+
+	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
+	if err != nil {
+		log.Printf("workload report dial error: %v", err)
+		return
+	}
+	defer client.Close()
+
+	base := func() *pb.ReportWorkloadsRequest {
+		return &pb.ReportWorkloadsRequest{
+			ServerId:     cfg.ServerID,
+			AgentToken:   cfg.AgentToken,
+			Hash:         hash,
+			DockerOk:     res.DockerOK,
+			DockerError:  res.DockerError,
+			SystemdOk:    res.SystemdOK,
+			SystemdError: res.SystemdError,
+		}
+	}
+
+	// The offer: hash only, no body. On an unchanged host this is the whole
+	// exchange, which is the point of the handshake.
+	needFull, err := client.ReportWorkloads(base())
+	if err != nil {
+		log.Printf("ReportWorkloads offer error: %v", err)
+		return
+	}
+	if !needFull {
+		return
+	}
+
+	req := base()
+	req.Full = true
+	req.Workloads = make([]pb.Workload, len(res.Workloads))
+	for i, w := range res.Workloads {
+		req.Workloads[i] = pb.Workload{
+			Kind:      w.Kind,
+			Id:        w.ID,
+			Name:      w.Name,
+			State:     w.State,
+			Health:    w.Health,
+			Image:     w.Image,
+			Stack:     w.Stack,
+			Ports:     w.Ports,
+			Restarts:  int32(w.Restarts),
+			Protected: w.Protected,
+		}
+		if !w.StartedAt.IsZero() {
+			req.Workloads[i].StartedAt = w.StartedAt.Format(time.RFC3339)
+		}
+	}
+
+	if _, err := client.ReportWorkloads(req); err != nil {
+		log.Printf("ReportWorkloads error: %v", err)
+		return
+	}
+	log.Printf("reported %d workload(s)", len(res.Workloads))
+}
+
+// handleRefreshWorkloads makes the agent report immediately. It sends nothing
+// back beyond the stream ack: the refresh is a nudge, not a channel, so there
+// is one writer for the collection rather than two.
+func handleRefreshWorkloads(cfg *config.Config) {
+	reportWorkloads(cfg)
+}
+
+// handleControlWorkload starts, stops or restarts a workload and answers with
+// the ordinary CommandResult.
+//
+// The agent's own protected check inside workloads.Control is the boundary; the
+// Protected flag it reports is only there so the UI can grey the button.
+func handleControlWorkload(send func(*pb.AgentMessage) error, cfg *config.Config, commandID string, cmd *pb.ControlWorkloadCmd) {
+	err := workloads.Control(context.Background(), cmd.Kind, cmd.Id, cmd.Action)
+
+	res := &pb.CommandResult{CommandId: commandID, Success: err == nil}
+	if err != nil {
+		res.Message = err.Error()
+		log.Printf("workload %s %s failed (cmd=%s): %v", cmd.Action, cmd.Id, commandID, err)
+	} else {
+		res.Message = cmd.Action + " " + cmd.Id + " ok"
+	}
+
+	_ = send(&pb.AgentMessage{
+		ServerId:   cfg.ServerID,
+		AgentToken: cfg.AgentToken,
+		Result:     res,
+	})
+
+	// Report straight away on success so the UI's refetch shows the new state
+	// rather than the old one.
+	if err == nil {
+		reportWorkloads(cfg)
+	}
+}
+
+func handleWorkloadLogs(send func(*pb.AgentMessage) error, cfg *config.Config, commandID string, cmd *pb.WorkloadLogsCmd) {
+	text, truncated, err := workloads.Logs(context.Background(), cmd.Kind, cmd.Id, int(cmd.Tail))
+	res := &pb.WorkloadLogsResult{CommandId: commandID, Text: text, Truncated: truncated}
+	if err != nil {
+		res.Error = err.Error()
+	}
+
+	_ = send(&pb.AgentMessage{
+		ServerId:           cfg.ServerID,
+		AgentToken:         cfg.AgentToken,
+		WorkloadLogsResult: res,
+	})
+}