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
@@ -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)
+	}
+}