fix: Fixes to running on kubernetes

2026-07-31 10:34:10 +01:00
parent de78688093
commit 165114471f
31 changed files with 1880 additions and 278 deletions
@@ -0,0 +1,220 @@
+// Package bus is the control plane's inter-process message bus.
+//
+// It exists because an agent's CommandStream lands on exactly one server
+// process. With a single replica that process is also the one handling every
+// REST request, so an in-memory map was enough. With several replicas it is
+// not: the pod asked to run a workflow step is almost never the pod holding
+// that agent's stream, and a map cannot reach across the gap.
+//
+// Redis is already a hard dependency (sessions), so the bus adds no new
+// infrastructure. It carries three things:
+//
+//	presence  which pod, if any, currently holds an agent's stream
+//	commands  a request/ack exchange delivering a ServerCommand to that pod
+//	results   step results travelling back to the pod driving the run
+//
+// Everything here is deliberately best-effort delivery with an explicit ack
+// rather than a queue. A command whose owner pod died between the presence
+// check and the publish must fail loudly and immediately — the caller answers
+// 503 and the operator retries — not sit in a queue waiting for a stream that
+// no longer exists.
+package bus
+
+import (
+	"context"
+	"crypto/rand"
+	"encoding/hex"
+	"encoding/json"
+	"errors"
+	"fmt"
+	"os"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+// ErrNoResponder means nothing was subscribed to the channel, or the subscriber
+// did not ack within the timeout. Either way the command did not reach an agent.
+var ErrNoResponder = errors.New("no responder on channel")
+
+var (
+	rdb    *redis.Client
+	nodeID string
+)
+
+// Init connects the bus. It takes its own client rather than sharing the
+// session store's: a subscription occupies its connection for as long as it
+// lives, and every agent stream on this pod holds one, so they must not come
+// out of the pool that ordinary session reads depend on.
+func Init(addr, username, password string) error {
+	rdb = redis.NewClient(&redis.Options{
+		Addr:     addr,
+		Username: username,
+		Password: password,
+	})
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	if err := rdb.Ping(ctx).Err(); err != nil {
+		return err
+	}
+
+	// Hostname is the pod name under Kubernetes, which makes a log line naming
+	// a node directly actionable. The random suffix keeps two processes on one
+	// host (or a reused pod name) from claiming each other's locks.
+	host, _ := os.Hostname()
+	if host == "" {
+		host = "server"
+	}
+	b := make([]byte, 4)
+	_, _ = rand.Read(b)
+	nodeID = host + "-" + hex.EncodeToString(b)
+	return nil
+}
+
+// NodeID identifies this process on the bus. Stable for the process lifetime.
+func NodeID() string { return nodeID }
+
+// Client exposes the underlying Redis client for the leader election and
+// presence helpers in this package. It is nil before Init.
+func Client() *redis.Client { return rdb }
+
+// Channel names. Every key and channel is prefixed so a Redis shared with the
+// session store (which uses km:) stays legible.
+const (
+	prefix = "vantage:"
+
+	// CommandChannel carries envelopes to whichever pod holds serverID's stream.
+	CommandChannel = prefix + "cmd:"
+	// ackChannel carries the owner pod's answer back to the requesting pod.
+	ackChannel = prefix + "ack:"
+	// ResultChannel carries a StepResult back to the pod driving the run.
+	ResultChannel = prefix + "res:"
+	// ProxyEndChannel carries a console relay's terminal reason back to the pod
+	// serving the WebSocket, which is the pod that has to write the audit event.
+	ProxyEndChannel = prefix + "proxyend:"
+
+	// PresenceKey records which node holds an agent's command stream.
+	PresenceKey = prefix + "agent:"
+	// leaderKey records the holder of a named singleton job.
+	leaderKey = prefix + "leader:"
+)
+
+// Publish sends v, JSON-encoded, to channel. It reports how many subscribers
+// received it, which is the only signal Redis pub/sub gives that anyone was
+// listening.
+func Publish(ctx context.Context, channel string, v any) (int64, error) {
+	b, err := json.Marshal(v)
+	if err != nil {
+		return 0, err
+	}
+	return rdb.Publish(ctx, channel, b).Result()
+}
+
+// Subscribe returns a channel of raw payloads and a function that unsubscribes.
+// It blocks until Redis confirms the subscription, so a caller may publish
+// immediately afterwards without racing its own subscriber.
+func Subscribe(ctx context.Context, channel string) (<-chan []byte, func(), error) {
+	ps := rdb.Subscribe(ctx, channel)
+	if _, err := ps.Receive(ctx); err != nil {
+		_ = ps.Close()
+		return nil, nil, err
+	}
+
+	out := make(chan []byte, 64)
+	go func() {
+		defer close(out)
+		for msg := range ps.Channel() {
+			select {
+			case out <- []byte(msg.Payload):
+			default:
+				// A subscriber too slow to keep up would otherwise stall every
+				// other subscriber sharing this connection. Dropping is correct
+				// here: commands are acked, and a dropped ack fails the caller
+				// loudly rather than hanging it.
+			}
+		}
+	}()
+
+	return out, func() { _ = ps.Close() }, nil
+}
+
+// Request publishes req to channel and waits for a single reply on replyTo.
+//
+// The reply subscription is established before the publish, so a responder that
+// answers instantly cannot beat the subscriber into place. A publish that
+// reaches no subscriber fails immediately with ErrNoResponder rather than
+// burning the whole timeout: nobody is going to answer.
+func Request(ctx context.Context, channel, replyTo string, req any, timeout time.Duration) ([]byte, error) {
+	waitCtx, cancel := context.WithTimeout(ctx, timeout)
+	defer cancel()
+
+	replies, unsub, err := Subscribe(waitCtx, replyTo)
+	if err != nil {
+		return nil, fmt.Errorf("subscribe reply channel: %w", err)
+	}
+	defer unsub()
+
+	n, err := Publish(waitCtx, channel, req)
+	if err != nil {
+		return nil, err
+	}
+	if n == 0 {
+		return nil, ErrNoResponder
+	}
+
+	select {
+	case b, ok := <-replies:
+		if !ok {
+			return nil, ErrNoResponder
+		}
+		return b, nil
+	case <-waitCtx.Done():
+		return nil, ErrNoResponder
+	}
+}
+
+// Reply answers a Request on its reply channel.
+func Reply(ctx context.Context, replyTo string, v any) error {
+	_, err := Publish(ctx, replyTo, v)
+	return err
+}
+
+// AckChannelFor names the reply channel for a command. The command ID is
+// already unique per dispatch, so it needs no further qualification.
+func AckChannelFor(commandID string) string { return ackChannel + commandID }
+
+// SetPresence claims serverID for this node for ttl. Called repeatedly by the
+// pod holding the stream; the TTL is what bounds how long a crashed pod keeps
+// claiming an agent it can no longer reach.
+func SetPresence(ctx context.Context, serverID string, ttl time.Duration) error {
+	return rdb.Set(ctx, PresenceKey+serverID, nodeID, ttl).Err()
+}
+
+// ClearPresence releases serverID, but only if this node still holds it. A
+// blind DEL would let a pod whose stream had already been re-established
+// elsewhere delete the new owner's claim on its way out.
+func ClearPresence(ctx context.Context, serverID string) error {
+	return releaseIfOwner.Run(ctx, rdb, []string{PresenceKey + serverID}, nodeID).Err()
+}
+
+// PresenceHolder returns the node holding serverID's stream, or "" if none does.
+func PresenceHolder(ctx context.Context, serverID string) string {
+	v, err := rdb.Get(ctx, PresenceKey+serverID).Result()
+	if err != nil {
+		return ""
+	}
+	return v
+}
+
+// IsConnected reports whether any pod currently holds serverID's stream.
+func IsConnected(ctx context.Context, serverID string) bool {
+	n, err := rdb.Exists(ctx, PresenceKey+serverID).Result()
+	return err == nil && n > 0
+}
+
+var releaseIfOwner = redis.NewScript(`
+if redis.call("GET", KEYS[1]) == ARGV[1] then
+	return redis.call("DEL", KEYS[1])
+end
+return 0
+`)
@@ -0,0 +1,100 @@
+package bus
+
+import (
+	"context"
+	"log"
+	"time"
+
+	"github.com/redis/go-redis/v9"
+)
+
+const (
+	leaderTTL   = 30 * time.Second
+	leaderRenew = 10 * time.Second
+)
+
+// RunAsLeader runs job on exactly one process at a time, cluster-wide.
+//
+// The background work this guards is not merely wasteful when duplicated. Every
+// replica running the monitor scheduler means every check fires N times, every
+// incident notification is delivered N times to the customer, and every hourly
+// rollup is written N times. The Free-instance reaper is worse: it deletes
+// whole instances, and two processes deleting the same one concurrently is not
+// a race anyone wins.
+//
+// Redis rather than a Kubernetes Lease so that Docker Compose, which has no
+// API server, takes the identical code path — one implementation to reason
+// about, not two.
+//
+// job is given a context cancelled the moment leadership is lost, and must
+// return when it is cancelled. Losing the lock (a paused process, a Redis
+// blip) is treated as fatal to that run of the job: the successor may already
+// have started, and two schedulers overlapping is the exact thing being
+// prevented.
+func RunAsLeader(ctx context.Context, name string, job func(context.Context)) {
+	go func() {
+		key := leaderKey + name
+		for {
+			if ctx.Err() != nil {
+				return
+			}
+			ok, err := rdb.SetNX(ctx, key, nodeID, leaderTTL).Result()
+			if err != nil {
+				log.Printf("leader %s: acquire failed: %v", name, err)
+			}
+			if !ok {
+				select {
+				case <-ctx.Done():
+					return
+				case <-time.After(leaderRenew):
+				}
+				continue
+			}
+
+			log.Printf("leader %s: acquired by %s", name, nodeID)
+			jobCtx, cancel := context.WithCancel(ctx)
+			go job(jobCtx)
+			holdLeadership(ctx, key, name)
+			cancel()
+			log.Printf("leader %s: released by %s", name, nodeID)
+
+			// Give up the key on a clean shutdown so a successor takes over in
+			// milliseconds rather than waiting out the TTL.
+			if ctx.Err() != nil {
+				relCtx, relCancel := context.WithTimeout(context.Background(), 2*time.Second)
+				_ = releaseIfOwner.Run(relCtx, rdb, []string{key}, nodeID).Err()
+				relCancel()
+				return
+			}
+		}
+	}()
+}
+
+// holdLeadership renews the lock until it is lost or ctx ends.
+func holdLeadership(ctx context.Context, key, name string) {
+	t := time.NewTicker(leaderRenew)
+	defer t.Stop()
+	for {
+		select {
+		case <-ctx.Done():
+			return
+		case <-t.C:
+			n, err := renewIfOwner.Run(ctx, rdb, []string{key}, nodeID, int(leaderTTL/time.Millisecond)).Int()
+			if err != nil && err != redis.Nil {
+				log.Printf("leader %s: renew failed, standing down: %v", name, err)
+				return
+			}
+			if n == 0 {
+				log.Printf("leader %s: lock lost, standing down", name)
+				return
+			}
+		}
+	}
+}
+
+var renewIfOwner = redis.NewScript(`
+if redis.call("GET", KEYS[1]) == ARGV[1] then
+	return redis.call("PEXPIRE", KEYS[1], ARGV[2])
+end
+return 0
+`)