feat: store workload reports and route log results

2026-08-07 08:56:26 +01:00
parent 6a4ef5b6c6
commit cf9d85b3cd
7 changed files with 403 additions and 15 deletions
@@ -40,6 +40,9 @@ type ReportWorkloadsRequest struct {
 	SystemdOk    bool       `json:"systemd_ok"`
 	SystemdError string     `json:"systemd_error,omitempty"`
 	Workloads    []Workload `json:"workloads,omitempty"` // empty on the offer call
+	// Full marks the second call. It is not inferred from an empty Workloads
+	// slice: a host running nothing sends an empty list as its full report.
+	Full bool `json:"full,omitempty"`
 }

 type ReportWorkloadsResponse struct {
@@ -163,6 +163,70 @@ func (s *vantageServer) ReportPackages(ctx context.Context, req *pb.ReportPackag
 	return &pb.ReportPackagesResponse{NeedFull: false}, nil
 }

+// ReportWorkloads stores what a server is running.
+//
+// It is not gated by licence: the workload registry reads as core fleet
+// management rather than a premium add-on. If that ever changes, the check
+// belongs here — gating collection, not display — for the same reason it does
+// in ReportPackages.
+func (s *vantageServer) ReportWorkloads(ctx context.Context, req *pb.ReportWorkloadsRequest) (*pb.ReportWorkloadsResponse, error) {
+	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
+	if err != nil {
+		return nil, status.Errorf(codes.Unauthenticated, "invalid agent token")
+	}
+
+	// The offer call: a hash and no body. Answering NeedFull=false here is what
+	// keeps an unchanged 60-second report to one small message.
+	//
+	// The offer is identified by Full, not by an empty Workloads slice: a host
+	// genuinely running nothing sends an empty list as its FULL report, and
+	// inferring the offer from emptiness would leave that host answering
+	// NeedFull=true forever and never storing anything.
+	if !req.Full {
+		known, err := services.HasWorkloadHash(srv.InstanceID, srv.ServerID, req.Hash)
+		if err != nil {
+			log.Printf("workload hash lookup for %s: %v", srv.ServerID, err)
+			return nil, status.Errorf(codes.Internal, "workload hash lookup failed")
+		}
+		return &pb.ReportWorkloadsResponse{NeedFull: !known}, nil
+	}
+
+	wls := make([]models.Workload, len(req.Workloads))
+	for i, w := range req.Workloads {
+		wls[i] = models.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:  int(w.Restarts),
+			Protected: w.Protected,
+		}
+		if w.StartedAt != "" {
+			if t, err := time.Parse(time.RFC3339, w.StartedAt); err == nil {
+				wls[i].StartedAt = t
+			}
+		}
+	}
+
+	if err := storeWorkloadReport(srv.InstanceID, srv.ServerID, req, wls); err != nil {
+		log.Printf("store workloads for %s: %v", srv.ServerID, err)
+		return nil, status.Errorf(codes.Internal, "failed to store workloads")
+	}
+	return &pb.ReportWorkloadsResponse{NeedFull: false}, nil
+}
+
+func storeWorkloadReport(instanceID, serverID string, req *pb.ReportWorkloadsRequest, wls []models.Workload) error {
+	if wls == nil {
+		wls = []models.Workload{}
+	}
+	return services.StoreWorkloads(instanceID, serverID, req.Hash, wls,
+		req.DockerOk, req.DockerError, req.SystemdOk, req.SystemdError)
+}
+
 func (s *vantageServer) ReportInventory(ctx context.Context, req *pb.InventoryReport) (*pb.InventoryReportResponse, error) {
 	srv, err := services.ValidateAgentToken(req.ServerId, req.AgentToken)
 	if err != nil {
@@ -257,6 +321,13 @@ func (s *vantageServer) CommandStream(stream pb.Vantage_CommandStreamServer) err
 			if m.Result != nil {
 				r := m.Result
 				log.Printf("agent %s cmd %s: success=%v %s", srv.ServerID, r.CommandId, r.Success, r.Message)
+				// Republished so a control action waiting on another pod sees
+				// it. Publishing with no subscriber is a no-op, so this is safe
+				// for every command result rather than only the awaited ones.
+				services.WorkloadResults.DeliverCommand(r)
+			}
+			if m.WorkloadLogsResult != nil {
+				services.WorkloadResults.Deliver(m.WorkloadLogsResult)
 			}
 			if m.StepResult != nil {
 				services.StepResults.Deliver(m.StepResult)
@@ -0,0 +1,128 @@
+package services
+
+import (
+	"context"
+	"encoding/json"
+	"log"
+
+	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/bus"
+	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/grpc/pb"
+)
+
+// Workload results travel back over the bus for the same reason commands travel
+// out over it: the pod serving the HTTP request and the pod holding the agent's
+// stream are two different processes, and a map in one cannot be read by the
+// other.
+//
+// Await MUST be called before the command is dispatched, or a fast agent
+// answers into a channel nobody is listening on yet. See stepresults.go.
+
+type workloadResultRegistry struct{}
+
+var WorkloadResults = &workloadResultRegistry{}
+
+// Await subscribes to a command's result channel for a log snapshot.
+func (r *workloadResultRegistry) Await(commandID string) (<-chan *pb.WorkloadLogsResult, func()) {
+	out := make(chan *pb.WorkloadLogsResult, 1)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	raw, unsub, err := bus.Subscribe(ctx, bus.ResultChannel+commandID)
+	if err != nil {
+		log.Printf("workload results: subscribe for %s: %v", commandID, err)
+		cancel()
+		close(out)
+		return out, func() {}
+	}
+
+	go func() {
+		defer close(out)
+		select {
+		case <-ctx.Done():
+			return
+		case b, ok := <-raw:
+			if !ok {
+				return
+			}
+			var res pb.WorkloadLogsResult
+			if err := json.Unmarshal(b, &res); err != nil {
+				log.Printf("workload results: undecodable result for %s: %v", commandID, err)
+				return
+			}
+			out <- &res
+		}
+	}()
+
+	return out, func() {
+		cancel()
+		unsub()
+	}
+}
+
+// AwaitCommand subscribes to a command's result channel for a plain
+// CommandResult, which is what a control action answers with.
+//
+// A control action reuses CommandResult rather than growing a message of its
+// own: start, stop and restart succeed or fail, and that is exactly what
+// CommandResult already says.
+func (r *workloadResultRegistry) AwaitCommand(commandID string) (<-chan *pb.CommandResult, func()) {
+	out := make(chan *pb.CommandResult, 1)
+
+	ctx, cancel := context.WithCancel(context.Background())
+	raw, unsub, err := bus.Subscribe(ctx, bus.ResultChannel+commandID)
+	if err != nil {
+		log.Printf("workload results: subscribe for %s: %v", commandID, err)
+		cancel()
+		close(out)
+		return out, func() {}
+	}
+
+	go func() {
+		defer close(out)
+		select {
+		case <-ctx.Done():
+			return
+		case b, ok := <-raw:
+			if !ok {
+				return
+			}
+			var res pb.CommandResult
+			if err := json.Unmarshal(b, &res); err != nil {
+				log.Printf("workload results: undecodable command result for %s: %v", commandID, err)
+				return
+			}
+			out <- &res
+		}
+	}()
+
+	return out, func() {
+		cancel()
+		unsub()
+	}
+}
+
+// Deliver publishes a log result received from an agent. Called on the pod
+// holding that agent's stream, which is not usually the pod waiting for it.
+func (r *workloadResultRegistry) Deliver(res *pb.WorkloadLogsResult) {
+	if res == nil || res.CommandId == "" {
+		return
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
+	defer cancel()
+	if _, err := bus.Publish(ctx, bus.ResultChannel+res.CommandId, res); err != nil {
+		log.Printf("workload results: publish for %s: %v", res.CommandId, err)
+	}
+}
+
+// DeliverCommand republishes a CommandResult onto the bus so a waiting pod can
+// see it. Publishing with no subscriber is a no-op, so this is safe to call for
+// every CommandResult rather than only the ones somebody is waiting on.
+func (r *workloadResultRegistry) DeliverCommand(res *pb.CommandResult) {
+	if res == nil || res.CommandId == "" {
+		return
+	}
+	ctx, cancel := context.WithTimeout(context.Background(), dispatchAckTimeout)
+	defer cancel()
+	if _, err := bus.Publish(ctx, bus.ResultChannel+res.CommandId, res); err != nil {
+		log.Printf("workload results: publish command result for %s: %v", res.CommandId, err)
+	}
+}
@@ -0,0 +1,187 @@
+package services
+
+import (
+	"context"
+	"fmt"
+	"time"
+
+	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/db"
+	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/grpc/pb"
+	"gitea.hostxtra.co.uk/mrhid6/vantage/server/internal/models"
+	"github.com/google/uuid"
+	"go.mongodb.org/mongo-driver/v2/bson"
+	"go.mongodb.org/mongo-driver/v2/mongo"
+	"go.mongodb.org/mongo-driver/v2/mongo/options"
+)
+
+// workloadResultTimeout bounds how long an API request waits for an agent to
+// answer a control action or a log read. It is well above the agent's own
+// 90-second control timeout and 60-second log timeout, so a slow-but-working
+// agent reports its real error rather than being cut off by this side.
+const workloadResultTimeout = 120 * time.Second
+
+func HasWorkloadHash(instanceID, serverID, hash string) (bool, error) {
+	err := db.Col("server_workloads").FindOne(context.Background(), bson.M{
+		"instance_id": instanceID,
+		"server_id":   serverID,
+		"hash":        hash,
+	}, options.FindOne().SetProjection(bson.M{"_id": 1})).Err()
+	if err == mongo.ErrNoDocuments {
+		return false, nil
+	}
+	return err == nil, err
+}
+
+// StoreWorkloads replaces a server's workload list.
+func StoreWorkloads(instanceID, serverID, hash string, wls []models.Workload,
+	dockerOK bool, dockerErr string, systemdOK bool, systemdErr string) error {
+
+	_, err := db.Col("server_workloads").UpdateOne(context.Background(),
+		bson.M{"instance_id": instanceID, "server_id": serverID},
+		bson.M{"$set": bson.M{
+			"hash":          hash,
+			"workloads":     wls,
+			"collected_at":  time.Now(),
+			"docker_ok":     dockerOK,
+			"docker_error":  dockerErr,
+			"systemd_ok":    systemdOK,
+			"systemd_error": systemdErr,
+		}},
+		options.UpdateOne().SetUpsert(true),
+	)
+	return err
+}
+
+func GetWorkloads(instanceID, serverID string) (*models.ServerWorkloads, error) {
+	var sw models.ServerWorkloads
+	err := db.Col("server_workloads").FindOne(context.Background(), bson.M{
+		"instance_id": instanceID,
+		"server_id":   serverID,
+	}).Decode(&sw)
+	if err == mongo.ErrNoDocuments {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+	return &sw, nil
+}
+
+type WorkloadHit struct {
+	ServerID string          `json:"server_id"`
+	Workload models.Workload `json:"workload"`
+}
+
+// SearchWorkloads answers "which servers run image X" — the reason the snapshot
+// is stored rather than fetched on demand and discarded.
+func SearchWorkloads(instanceID, image, stack, state string) ([]WorkloadHit, error) {
+	ctx := context.Background()
+
+	filter := bson.M{"instance_id": instanceID}
+	if image != "" {
+		filter["workloads.image"] = image
+	}
+
+	cur, err := db.Col("server_workloads").Find(ctx, filter)
+	if err != nil {
+		return nil, err
+	}
+	defer cur.Close(ctx)
+
+	var docs []models.ServerWorkloads
+	if err := cur.All(ctx, &docs); err != nil {
+		return nil, err
+	}
+
+	hits := []WorkloadHit{}
+	for _, d := range docs {
+		for _, w := range d.Workloads {
+			if image != "" && w.Image != image {
+				continue
+			}
+			if stack != "" && w.Stack != stack {
+				continue
+			}
+			if state != "" && w.State != state {
+				continue
+			}
+			hits = append(hits, WorkloadHit{ServerID: d.ServerID, Workload: w})
+		}
+	}
+	return hits, nil
+}
+
+// DispatchRefreshWorkloads asks an agent to report immediately. It returns as
+// soon as the owning pod acks; the caller refetches the stored document.
+//
+// The refresh carries nothing back on purpose: the agent answers through the
+// normal ReportWorkloads RPC, so server_workloads has exactly one writer.
+func DispatchRefreshWorkloads(serverID string) error {
+	return Dispatcher.dispatch(serverID, &pb.ServerCommand{
+		CommandId:        uuid.New().String(),
+		RefreshWorkloads: &pb.RefreshWorkloadsCmd{},
+	})
+}
+
+// DispatchControlWorkload runs a control action and waits for the agent's
+// CommandResult.
+//
+// Await is called BEFORE dispatch. Reversing those two lines introduces a race
+// that only shows under load, on a fast agent answering into a channel nobody
+// has joined yet.
+func DispatchControlWorkload(serverID, kind, id, action string) error {
+	commandID := uuid.New().String()
+
+	results, done := WorkloadResults.AwaitCommand(commandID)
+	defer done()
+
+	if err := Dispatcher.dispatch(serverID, &pb.ServerCommand{
+		CommandId:       commandID,
+		ControlWorkload: &pb.ControlWorkloadCmd{Kind: kind, Id: id, Action: action},
+	}); err != nil {
+		return err
+	}
+
+	select {
+	case res, ok := <-results:
+		if !ok || res == nil {
+			return fmt.Errorf("no result from agent for %s %s", action, id)
+		}
+		if !res.Success {
+			return fmt.Errorf("%s", res.Message)
+		}
+		return nil
+	case <-time.After(workloadResultTimeout):
+		return fmt.Errorf("timed out waiting for the agent to %s %s", action, id)
+	}
+}
+
+// DispatchWorkloadLogs fetches a bounded log snapshot.
+//
+// Await is called BEFORE dispatch, for the same reason as above.
+func DispatchWorkloadLogs(serverID, kind, id string, tail int) (string, bool, error) {
+	commandID := uuid.New().String()
+
+	results, done := WorkloadResults.Await(commandID)
+	defer done()
+
+	if err := Dispatcher.dispatch(serverID, &pb.ServerCommand{
+		CommandId:    commandID,
+		WorkloadLogs: &pb.WorkloadLogsCmd{Kind: kind, Id: id, Tail: int32(tail)},
+	}); err != nil {
+		return "", false, err
+	}
+
+	select {
+	case res, ok := <-results:
+		if !ok || res == nil {
+			return "", false, fmt.Errorf("no log result from agent for %s", id)
+		}
+		if res.Error != "" {
+			return "", false, fmt.Errorf("%s", res.Error)
+		}
+		return res.Text, res.Truncated, nil
+	case <-time.After(workloadResultTimeout):
+		return "", false, fmt.Errorf("timed out waiting for logs for %s", id)
+	}
+}