Cluster heartbeat support, improved cluster start/stop, coordinated shutdown in multi-server mode

2026-10-21 11:46:37 +00:00 · 2026-02-14 16:59:32 +10:00
parent c5e561abbf
commit 24924be941
11 changed files with 1088 additions and 6 deletions
@@ -75,6 +75,21 @@ var adminCmd = &cobra.Command{
 	Long:  `Administrative commands for initializing and managing the Hive platform.`,
 }

+var clusterCmd = &cobra.Command{
+	Use:   "cluster",
+	Short: "Cluster-wide operations",
+	Long:  `Cluster-wide administrative operations such as coordinated shutdown.`,
+}
+
+var clusterShutdownCmd = &cobra.Command{
+	Use:   "shutdown",
+	Short: "Gracefully shut down the entire cluster",
+	Long: `Perform a coordinated, phased shutdown of the entire cluster.
+Phases execute in order: GATE (stop API/UI) → DRAIN (stop VMs) → STORAGE (stop viperblock) → PERSIST (stop predastore) → INFRA (stop NATS/daemon).
+Each phase waits for all nodes to ACK before proceeding to the next.`,
+	Run: runClusterShutdown,
+}
+
 var adminInitCmd = &cobra.Command{
 	Use:   "init",
 	Short: "Initialize Hive platform configuration",
@@ -137,6 +152,12 @@ func init() {
 	adminCmd.AddCommand(adminInitCmd)
 	adminCmd.AddCommand(adminJoinCmd)

+	adminCmd.AddCommand(clusterCmd)
+	clusterCmd.AddCommand(clusterShutdownCmd)
+	clusterShutdownCmd.Flags().Bool("force", false, "Force shutdown even if nodes don't respond")
+	clusterShutdownCmd.Flags().Duration("timeout", 120*time.Second, "Maximum time to wait per phase")
+	clusterShutdownCmd.Flags().Bool("dry-run", false, "Print phase plan without executing")
+
 	adminCmd.AddCommand(imagesCmd)
 	imagesCmd.AddCommand(imagesImportCmd)
 	imagesCmd.AddCommand(imagesListCmd)
@@ -0,0 +1,195 @@
+package cmd
+
+import (
+	"encoding/json"
+	"fmt"
+	"os"
+	"strings"
+	"time"
+
+	"github.com/mulgadc/hive/hive/daemon"
+	"github.com/nats-io/nats.go"
+	"github.com/spf13/cobra"
+)
+
+// runClusterShutdown orchestrates a phased, coordinated shutdown of the cluster.
+func runClusterShutdown(cmd *cobra.Command, args []string) {
+	force, _ := cmd.Flags().GetBool("force")
+	timeout, _ := cmd.Flags().GetDuration("timeout")
+	dryRun, _ := cmd.Flags().GetBool("dry-run")
+
+	cfg, nc, err := loadConfigAndConnect()
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
+		os.Exit(1)
+	}
+	defer nc.Close()
+
+	nodeCount := len(cfg.Nodes)
+	phases := []string{"gate", "drain", "storage", "persist", "infra"}
+
+	if dryRun {
+		fmt.Println("Cluster Shutdown Plan (dry-run)")
+		fmt.Printf("  Nodes: %d\n", nodeCount)
+		for name, nodeCfg := range cfg.Nodes {
+			fmt.Printf("    - %s (services: %s)\n", name, strings.Join(nodeCfg.GetServices(), ", "))
+		}
+		fmt.Printf("  Phases: %s\n", strings.Join(phases, " -> "))
+		fmt.Printf("  Timeout per phase: %s\n", timeout)
+		fmt.Printf("  Force: %v\n", force)
+		fmt.Println("\nPhase details:")
+		fmt.Println("  1. GATE    - Stop API gateway and UI, reject new work")
+		fmt.Println("  2. DRAIN   - Gracefully stop all VMs, persist state")
+		fmt.Println("  3. STORAGE - Stop viperblock, cleanup nbdkit")
+		fmt.Println("  4. PERSIST - Stop predastore")
+		fmt.Println("  5. INFRA   - Stop NATS, exit daemons")
+		return
+	}
+
+	fmt.Printf("Starting coordinated cluster shutdown (%d nodes)\n", nodeCount)
+	fmt.Printf("Phases: %s\n", strings.Join(phases, " -> "))
+	fmt.Printf("Timeout per phase: %s\n\n", timeout)
+
+	// Write cluster shutdown marker to KV
+	jsm, err := daemon.NewJetStreamManager(nc, 1)
+	if err != nil {
+		fmt.Fprintf(os.Stderr, "Warning: failed to create JetStream manager: %v\n", err)
+	} else {
+		if err := jsm.InitClusterStateBucket(); err != nil {
+			fmt.Fprintf(os.Stderr, "Warning: failed to init cluster state bucket: %v\n", err)
+		} else {
+			state := &daemon.ClusterShutdownState{
+				Initiator:  cfg.Node,
+				Phase:      "starting",
+				Started:    time.Now().UTC().Format(time.RFC3339),
+				Timeout:    timeout.String(),
+				Force:      force,
+				NodesTotal: nodeCount,
+				NodesAcked: make(map[string]string),
+			}
+			if err := jsm.WriteClusterShutdown(state); err != nil {
+				fmt.Fprintf(os.Stderr, "Warning: failed to write shutdown state: %v\n", err)
+			}
+		}
+	}
+
+	start := time.Now()
+
+	// Execute phases sequentially (except INFRA which is fire-and-forget)
+	for _, phase := range phases {
+		topic := "hive.cluster.shutdown." + phase
+		req := daemon.ShutdownRequest{
+			Phase:   phase,
+			Force:   force,
+			Timeout: int(timeout.Seconds()),
+		}
+
+		reqData, err := json.Marshal(req)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "Error marshaling request: %v\n", err)
+			os.Exit(1)
+		}
+
+		if phase == "infra" {
+			// INFRA is fire-and-forget — NATS is going down, no ACKs possible
+			fmt.Printf("[INFRA] Sending final shutdown to all nodes...\n")
+			if err := nc.Publish(topic, reqData); err != nil {
+				fmt.Fprintf(os.Stderr, "Error publishing infra shutdown: %v\n", err)
+			}
+			nc.Flush()
+			// Wait briefly for messages to propagate
+			time.Sleep(2 * time.Second)
+			fmt.Printf("[INFRA] Complete\n")
+			break
+		}
+
+		// For DRAIN phase, subscribe to progress updates
+		var progressSub *nats.Subscription
+		if phase == "drain" {
+			progressSub, err = nc.Subscribe("hive.cluster.shutdown.progress", func(msg *nats.Msg) {
+				var progress daemon.ShutdownProgress
+				if err := json.Unmarshal(msg.Data, &progress); err != nil {
+					return
+				}
+				fmt.Printf("  [%s] %s: %d/%d VMs remaining\n", strings.ToUpper(phase), progress.Node, progress.Remaining, progress.Total)
+			})
+			if err != nil {
+				fmt.Fprintf(os.Stderr, "Warning: failed to subscribe to progress: %v\n", err)
+			}
+		}
+
+		// Collect ACKs from all nodes
+		phaseStart := time.Now()
+		fmt.Printf("[%s] Sending to %d node(s)...\n", strings.ToUpper(phase), nodeCount)
+
+		acks, err := collectShutdownACKs(nc, topic, reqData, nodeCount, timeout)
+		if err != nil {
+			fmt.Fprintf(os.Stderr, "[%s] Error: %v\n", strings.ToUpper(phase), err)
+			if !force {
+				fmt.Fprintf(os.Stderr, "Aborting shutdown. Use --force to continue despite errors.\n")
+				os.Exit(1)
+			}
+		}
+
+		// Unsubscribe from progress
+		if progressSub != nil {
+			progressSub.Unsubscribe()
+		}
+
+		// Print results
+		for _, ack := range acks {
+			if ack.Error != "" {
+				fmt.Printf("  [%s] %s: ERROR - %s\n", strings.ToUpper(phase), ack.Node, ack.Error)
+			} else if len(ack.Stopped) > 0 {
+				fmt.Printf("  [%s] %s: stopped %s\n", strings.ToUpper(phase), ack.Node, strings.Join(ack.Stopped, ", "))
+			} else {
+				fmt.Printf("  [%s] %s: OK\n", strings.ToUpper(phase), ack.Node)
+			}
+		}
+
+		ackedCount := len(acks)
+		if ackedCount < nodeCount && !force {
+			fmt.Fprintf(os.Stderr, "[%s] Only %d/%d nodes responded. Use --force to continue.\n",
+				strings.ToUpper(phase), ackedCount, nodeCount)
+			os.Exit(1)
+		}
+
+		fmt.Printf("[%s] Complete (%d/%d nodes, %s)\n\n", strings.ToUpper(phase), ackedCount, nodeCount, time.Since(phaseStart).Round(time.Millisecond))
+	}
+
+	fmt.Printf("Cluster shutdown complete (%s)\n", time.Since(start).Round(time.Millisecond))
+}
+
+// collectShutdownACKs publishes a shutdown request and collects ACKs from nodes.
+func collectShutdownACKs(nc *nats.Conn, topic string, reqData []byte, nodeCount int, timeout time.Duration) ([]daemon.ShutdownACK, error) {
+	inbox := nats.NewInbox()
+	sub, err := nc.SubscribeSync(inbox)
+	if err != nil {
+		return nil, fmt.Errorf("failed to subscribe to inbox: %w", err)
+	}
+	defer sub.Unsubscribe()
+
+	if err := nc.PublishRequest(topic, inbox, reqData); err != nil {
+		return nil, fmt.Errorf("failed to publish request: %w", err)
+	}
+	nc.Flush()
+
+	var acks []daemon.ShutdownACK
+	deadline := time.Now().Add(timeout)
+	for len(acks) < nodeCount {
+		remaining := time.Until(deadline)
+		if remaining <= 0 {
+			break
+		}
+		msg, err := sub.NextMsg(remaining)
+		if err != nil {
+			break
+		}
+		var ack daemon.ShutdownACK
+		if err := json.Unmarshal(msg.Data, &ack); err != nil {
+			continue
+		}
+		acks = append(acks, ack)
+	}
+	return acks, nil
+}