mirror of
https://github.com/mulgadc/spinifex.git
synced 2026-10-15 09:01:31 +00:00
Cluster heartbeat support, improved cluster start/stop, coordinated shutdown in multi-server mode
This commit is contained in:
Ben Duncan
@@ -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
|
||||
}
|
||||
@@ -0,0 +1,78 @@
|
||||
# Phase 2.1 Heartbeats + Phase 2.4 Coordinated Cluster Shutdown
|
||||
|
||||
**Status: Complete**
|
||||
|
||||
## Summary
|
||||
|
||||
Implemented daemon heartbeats (Phase 2.1) and coordinated cluster shutdown (Phase 2.4). Heartbeats publish node status to KV every 10s. Coordinated shutdown drives a phased shutdown across all nodes via NATS, replacing the error-prone per-node approach for multi-node clusters.
|
||||
|
||||
## Context / Problem Statement
|
||||
|
||||
- No cluster-wide health awareness — nodes don't know peer status
|
||||
- Multi-node shutdown required manual per-node `stop-dev.sh`, which is error-prone and doesn't guarantee correct phase ordering (VMs must stop before storage, storage before NATS)
|
||||
- Phase 0+1 (capabilities, recovery throttling, shutdown markers) were already complete
|
||||
|
||||
## Design Decisions
|
||||
|
||||
- **Option B chosen**: Daemons handle stopping local services internally via `utils.StopProcess()` when they receive NATS shutdown phase messages. The CLI coordinator just sends messages and collects ACKs.
|
||||
- **Heartbeat interval**: 10 seconds — fast enough for health monitoring, low overhead
|
||||
- **Shutdown phases**: GATE → DRAIN → STORAGE → PERSIST → INFRA — ordered to prevent data loss
|
||||
- **INFRA phase is fire-and-forget**: NATS is going down, no ACK possible
|
||||
- **Scripts**: `stop-dev.sh` auto-detects multi-node and delegates to `hive admin cluster shutdown`; `HIVE_FORCE_LOCAL_STOP=1` overrides for per-service stop
|
||||
|
||||
## Proposed Changes (Implemented)
|
||||
|
||||
### Heartbeat System
|
||||
- `Heartbeat` struct in `jetstream.go` with node, epoch, services, VM count, resource allocation
|
||||
- `WriteHeartbeat`/`ReadHeartbeat` KV methods on `JetStreamManager`
|
||||
- `startHeartbeat()` goroutine in `heartbeat.go` — fires immediately, then every 10s
|
||||
- `buildHeartbeat()` reads from `ResourceManager` and `Instances`
|
||||
|
||||
### Coordinated Shutdown
|
||||
- Five shutdown phase handlers in `shutdown.go`: GATE, DRAIN, STORAGE, PERSIST, INFRA
|
||||
- `ShutdownRequest`, `ShutdownACK`, `ShutdownProgress` message types
|
||||
- `shuttingDown atomic.Bool` on Daemon — set during GATE, checked by `setupShutdown()` to skip redundant VM stops
|
||||
- NATS subscriptions for `hive.cluster.shutdown.{phase}` (fan-out, no queue group)
|
||||
- Progress updates published during DRAIN phase
|
||||
|
||||
### CLI Coordinator
|
||||
- `hive admin cluster shutdown` command in `cluster.go`
|
||||
- Flags: `--force`, `--timeout`, `--dry-run`
|
||||
- Executes phases sequentially, collecting ACKs with timeout
|
||||
- Writes `ClusterShutdownState` to KV for observability
|
||||
- Live progress display during DRAIN phase
|
||||
|
||||
### Script Updates
|
||||
- `stop-dev.sh`: `is_multinode()` detection, delegates to coordinated shutdown
|
||||
- `start-dev.sh`: `is_multinode()` detection, cluster health poll at end
|
||||
|
||||
## Files Modified
|
||||
|
||||
| File | Action | Description |
|
||||
|------|--------|-------------|
|
||||
| `hive/daemon/jetstream.go` | Edit | `Heartbeat`, `ClusterShutdownState` structs + KV methods |
|
||||
| `hive/daemon/heartbeat.go` | New | `startHeartbeat()`, `buildHeartbeat()` goroutine |
|
||||
| `hive/daemon/shutdown.go` | New | Shutdown types + 5 phase handlers |
|
||||
| `hive/daemon/daemon.go` | Edit | `shuttingDown atomic.Bool`, shutdown subscriptions, `startHeartbeat()`, updated `setupShutdown()` |
|
||||
| `cmd/hive/cmd/admin.go` | Edit | `clusterCmd` + `clusterShutdownCmd` registration |
|
||||
| `cmd/hive/cmd/cluster.go` | New | `runClusterShutdown()` coordinator |
|
||||
| `scripts/stop-dev.sh` | Edit | Multi-node detection + coordinated shutdown delegation |
|
||||
| `scripts/start-dev.sh` | Edit | Multi-node detection + cluster health poll |
|
||||
| `hive/daemon/heartbeat_test.go` | New | Heartbeat tests (build, allocation, KV round-trip) |
|
||||
| `hive/daemon/shutdown_test.go` | New | Shutdown tests (marshal, KV round-trip, flag behavior) |
|
||||
|
||||
## Testing
|
||||
|
||||
- `make preflight` — all checks pass (gofmt, go vet, gosec, staticcheck, govulncheck, all tests)
|
||||
- 9 new unit tests covering heartbeat building, resource reflection, KV round-trips, shutdown message marshaling, and shutdown state management
|
||||
- `hive admin cluster shutdown --dry-run` prints phase plan
|
||||
- `stop-dev.sh` on multi-node delegates to coordinated shutdown
|
||||
- `HIVE_FORCE_LOCAL_STOP=1 stop-dev.sh` on multi-node uses per-service stop
|
||||
- `stop-dev.sh` on single-node uses per-service stop (unchanged behavior)
|
||||
|
||||
## Future Work
|
||||
|
||||
- Phase 2.2: Heartbeat-driven health monitoring (stale heartbeat detection, node failure alerts)
|
||||
- Phase 2.3: Automatic VM migration on node failure
|
||||
- Integration with `hive top` to display heartbeat data
|
||||
- Cluster shutdown progress bar in UI
|
||||
+21
-6
@@ -19,6 +19,7 @@ import (
|
||||
"strconv"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"syscall"
|
||||
"time"
|
||||
|
||||
@@ -117,6 +118,10 @@ type Daemon struct {
|
||||
// Delay after QMP device_del before blockdev-del (default 1s, 0 in tests)
|
||||
detachDelay time.Duration
|
||||
|
||||
// shuttingDown is set to true during coordinated cluster shutdown (GATE phase).
|
||||
// When true, the daemon rejects new work and setupShutdown skips VM stops.
|
||||
shuttingDown atomic.Bool
|
||||
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
@@ -529,6 +534,12 @@ func (d *Daemon) subscribeAll() error {
|
||||
{"hive.nodes.discover", d.handleNodeDiscover, ""},
|
||||
{"hive.node.status", d.handleNodeStatus, ""},
|
||||
{"hive.node.vms", d.handleNodeVMs, ""},
|
||||
// Coordinated cluster shutdown phases (fan-out, no queue group)
|
||||
{"hive.cluster.shutdown.gate", d.handleShutdownGate, ""},
|
||||
{"hive.cluster.shutdown.drain", d.handleShutdownDrain, ""},
|
||||
{"hive.cluster.shutdown.storage", d.handleShutdownStorage, ""},
|
||||
{"hive.cluster.shutdown.persist", d.handleShutdownPersist, ""},
|
||||
{"hive.cluster.shutdown.infra", d.handleShutdownInfra, ""},
|
||||
}
|
||||
|
||||
for _, s := range subs {
|
||||
@@ -610,6 +621,7 @@ func (d *Daemon) Start() error {
|
||||
// are only routed to nodes with available capacity.
|
||||
d.resourceMgr.initSubscriptions(d.natsConn, d.handleEC2RunInstances)
|
||||
|
||||
d.startHeartbeat()
|
||||
d.setupShutdown()
|
||||
d.awaitShutdown()
|
||||
|
||||
@@ -1378,12 +1390,17 @@ func (d *Daemon) setupShutdown() {
|
||||
<-sigChan
|
||||
slog.Info("Received shutdown signal, cleaning up...")
|
||||
|
||||
// Cancel context
|
||||
// Cancel context to stop heartbeat and other goroutines
|
||||
d.cancel()
|
||||
|
||||
// Pass instances to terminate
|
||||
if err := d.stopInstance(d.Instances.VMS, false); err != nil {
|
||||
slog.Error("Failed to stop instances during shutdown", "err", err)
|
||||
// If coordinated shutdown already handled VMs (DRAIN phase), skip stopInstance
|
||||
if d.shuttingDown.Load() {
|
||||
slog.Info("Coordinated shutdown in progress, skipping VM stop (already handled by DRAIN phase)")
|
||||
} else {
|
||||
// Pass instances to terminate
|
||||
if err := d.stopInstance(d.Instances.VMS, false); err != nil {
|
||||
slog.Error("Failed to stop instances during shutdown", "err", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Final cleanup
|
||||
@@ -1420,8 +1437,6 @@ func (d *Daemon) setupShutdown() {
|
||||
}
|
||||
}
|
||||
|
||||
// Wait for any ongoing operations to complete
|
||||
// TODO: Implement cleanup of running instances
|
||||
slog.Info("Shutdown complete")
|
||||
})
|
||||
}
|
||||
|
||||
@@ -0,0 +1,69 @@
|
||||
package daemon
|
||||
|
||||
import (
|
||||
"log/slog"
|
||||
"time"
|
||||
)
|
||||
|
||||
const heartbeatInterval = 10 * time.Second
|
||||
|
||||
// startHeartbeat launches a goroutine that publishes this daemon's heartbeat
|
||||
// to the cluster-state KV store every heartbeatInterval. It fires immediately
|
||||
// on start, then repeats on a ticker. The goroutine exits when d.ctx is cancelled.
|
||||
func (d *Daemon) startHeartbeat() {
|
||||
if d.jsManager == nil {
|
||||
slog.Warn("JetStream not initialized, skipping heartbeat")
|
||||
return
|
||||
}
|
||||
|
||||
go func() {
|
||||
// Fire immediately on startup
|
||||
d.publishHeartbeat()
|
||||
|
||||
ticker := time.NewTicker(heartbeatInterval)
|
||||
defer ticker.Stop()
|
||||
|
||||
for {
|
||||
select {
|
||||
case <-d.ctx.Done():
|
||||
slog.Info("Heartbeat goroutine stopping")
|
||||
return
|
||||
case <-ticker.C:
|
||||
d.publishHeartbeat()
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
slog.Info("Heartbeat started", "interval", heartbeatInterval)
|
||||
}
|
||||
|
||||
// publishHeartbeat builds and writes a heartbeat entry to KV.
|
||||
func (d *Daemon) publishHeartbeat() {
|
||||
h := d.buildHeartbeat()
|
||||
if err := d.jsManager.WriteHeartbeat(h); err != nil {
|
||||
slog.Warn("Failed to publish heartbeat", "error", err)
|
||||
} else {
|
||||
slog.Debug("Heartbeat published", "node", h.Node, "vms", h.VMCount)
|
||||
}
|
||||
}
|
||||
|
||||
// buildHeartbeat constructs a Heartbeat from the daemon's current state.
|
||||
func (d *Daemon) buildHeartbeat() *Heartbeat {
|
||||
totalVCPU, totalMem, allocVCPU, allocMem, _ := d.resourceMgr.GetResourceStats()
|
||||
|
||||
d.Instances.Mu.Lock()
|
||||
vmCount := len(d.Instances.VMS)
|
||||
d.Instances.Mu.Unlock()
|
||||
|
||||
return &Heartbeat{
|
||||
Node: d.node,
|
||||
Epoch: d.clusterConfig.Epoch,
|
||||
Timestamp: time.Now().UTC().Format(time.RFC3339),
|
||||
Services: d.config.GetServices(),
|
||||
VMCount: vmCount,
|
||||
AllocatedVCPU: allocVCPU,
|
||||
AvailableVCPU: totalVCPU - allocVCPU,
|
||||
AllocatedMem: allocMem,
|
||||
AvailableMem: totalMem - allocMem,
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,126 @@
|
||||
package daemon
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/mulgadc/hive/hive/config"
|
||||
"github.com/mulgadc/hive/hive/vm"
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestBuildHeartbeat verifies that buildHeartbeat populates the struct from daemon state.
|
||||
func TestBuildHeartbeat(t *testing.T) {
|
||||
d := &Daemon{
|
||||
node: "test-node",
|
||||
clusterConfig: &config.ClusterConfig{
|
||||
Epoch: 5,
|
||||
},
|
||||
config: &config.Config{
|
||||
Services: []string{"daemon", "nats", "viperblock"},
|
||||
},
|
||||
resourceMgr: NewResourceManager(),
|
||||
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
|
||||
}
|
||||
|
||||
h := d.buildHeartbeat()
|
||||
|
||||
assert.Equal(t, "test-node", h.Node)
|
||||
assert.Equal(t, uint64(5), h.Epoch)
|
||||
assert.NotEmpty(t, h.Timestamp)
|
||||
assert.Equal(t, []string{"daemon", "nats", "viperblock"}, h.Services)
|
||||
assert.Equal(t, 0, h.VMCount)
|
||||
assert.Equal(t, 0, h.AllocatedVCPU)
|
||||
assert.Greater(t, h.AvailableVCPU, 0)
|
||||
assert.Greater(t, h.AvailableMem, 0.0)
|
||||
}
|
||||
|
||||
// TestHeartbeatReflectsAllocation verifies that allocating resources changes the heartbeat values.
|
||||
func TestHeartbeatReflectsAllocation(t *testing.T) {
|
||||
d := &Daemon{
|
||||
node: "alloc-node",
|
||||
clusterConfig: &config.ClusterConfig{
|
||||
Epoch: 1,
|
||||
},
|
||||
config: &config.Config{
|
||||
Services: []string{"daemon"},
|
||||
},
|
||||
resourceMgr: NewResourceManager(),
|
||||
Instances: vm.Instances{VMS: make(map[string]*vm.VM)},
|
||||
}
|
||||
|
||||
// Take a heartbeat before allocation
|
||||
before := d.buildHeartbeat()
|
||||
|
||||
// Find an instance type we can allocate
|
||||
var allocType string
|
||||
for typeName := range d.resourceMgr.instanceTypes {
|
||||
if d.resourceMgr.canAllocate(d.resourceMgr.instanceTypes[typeName], 1) >= 1 {
|
||||
allocType = typeName
|
||||
break
|
||||
}
|
||||
}
|
||||
require.NotEmpty(t, allocType, "Should have at least one allocatable instance type")
|
||||
|
||||
// Allocate resources
|
||||
err := d.resourceMgr.allocate(d.resourceMgr.instanceTypes[allocType])
|
||||
require.NoError(t, err)
|
||||
|
||||
// Add a VM to the instance map
|
||||
d.Instances.Mu.Lock()
|
||||
d.Instances.VMS["i-test-001"] = &vm.VM{
|
||||
ID: "i-test-001",
|
||||
Status: vm.StateRunning,
|
||||
InstanceType: allocType,
|
||||
}
|
||||
d.Instances.Mu.Unlock()
|
||||
|
||||
// Take a heartbeat after allocation
|
||||
after := d.buildHeartbeat()
|
||||
|
||||
assert.Equal(t, 1, after.VMCount, "Should reflect 1 VM")
|
||||
assert.Greater(t, after.AllocatedVCPU, before.AllocatedVCPU, "AllocatedVCPU should increase")
|
||||
assert.Less(t, after.AvailableVCPU, before.AvailableVCPU, "AvailableVCPU should decrease")
|
||||
}
|
||||
|
||||
// TestHeartbeatKVRoundTrip verifies that heartbeat can be written to and read from KV.
|
||||
func TestHeartbeatKVRoundTrip(t *testing.T) {
|
||||
nc, err := nats.Connect(sharedJSNATSURL)
|
||||
require.NoError(t, err)
|
||||
defer nc.Close()
|
||||
|
||||
jsm, err := NewJetStreamManager(nc, 1)
|
||||
require.NoError(t, err)
|
||||
err = jsm.InitClusterStateBucket()
|
||||
require.NoError(t, err)
|
||||
|
||||
h := &Heartbeat{
|
||||
Node: "kv-test-node",
|
||||
Epoch: 3,
|
||||
Timestamp: "2025-01-01T00:00:00Z",
|
||||
Services: []string{"daemon", "nats"},
|
||||
VMCount: 2,
|
||||
AllocatedVCPU: 4,
|
||||
AvailableVCPU: 12,
|
||||
AllocatedMem: 8.0,
|
||||
AvailableMem: 24.0,
|
||||
}
|
||||
|
||||
err = jsm.WriteHeartbeat(h)
|
||||
require.NoError(t, err)
|
||||
|
||||
loaded, err := jsm.ReadHeartbeat("kv-test-node")
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, loaded)
|
||||
|
||||
assert.Equal(t, h.Node, loaded.Node)
|
||||
assert.Equal(t, h.Epoch, loaded.Epoch)
|
||||
assert.Equal(t, h.Timestamp, loaded.Timestamp)
|
||||
assert.Equal(t, h.Services, loaded.Services)
|
||||
assert.Equal(t, h.VMCount, loaded.VMCount)
|
||||
assert.Equal(t, h.AllocatedVCPU, loaded.AllocatedVCPU)
|
||||
assert.Equal(t, h.AvailableVCPU, loaded.AvailableVCPU)
|
||||
assert.Equal(t, h.AllocatedMem, loaded.AllocatedMem)
|
||||
assert.Equal(t, h.AvailableMem, loaded.AvailableMem)
|
||||
}
|
||||
@@ -99,6 +99,100 @@ func (m *JetStreamManager) InitClusterStateBucket() error {
|
||||
return nil
|
||||
}
|
||||
|
||||
// Heartbeat represents a daemon's periodic health status published to cluster KV.
|
||||
type Heartbeat struct {
|
||||
Node string `json:"node"`
|
||||
Epoch uint64 `json:"epoch"`
|
||||
Timestamp string `json:"timestamp"`
|
||||
Services []string `json:"services"`
|
||||
VMCount int `json:"vm_count"`
|
||||
AllocatedVCPU int `json:"allocated_vcpu"`
|
||||
AvailableVCPU int `json:"available_vcpu"`
|
||||
AllocatedMem float64 `json:"allocated_mem_gb"`
|
||||
AvailableMem float64 `json:"available_mem_gb"`
|
||||
}
|
||||
|
||||
// WriteHeartbeat writes a heartbeat entry for the given node to the cluster-state KV.
|
||||
func (m *JetStreamManager) WriteHeartbeat(h *Heartbeat) error {
|
||||
if m.clusterKV == nil {
|
||||
return errors.New("cluster state KV not initialized")
|
||||
}
|
||||
data, err := json.Marshal(h)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = m.clusterKV.Put("heartbeat."+h.Node, data)
|
||||
return err
|
||||
}
|
||||
|
||||
// ReadHeartbeat reads the heartbeat entry for the given node from the cluster-state KV.
|
||||
func (m *JetStreamManager) ReadHeartbeat(nodeID string) (*Heartbeat, error) {
|
||||
if m.clusterKV == nil {
|
||||
return nil, errors.New("cluster state KV not initialized")
|
||||
}
|
||||
entry, err := m.clusterKV.Get("heartbeat." + nodeID)
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var h Heartbeat
|
||||
if err := json.Unmarshal(entry.Value(), &h); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &h, nil
|
||||
}
|
||||
|
||||
// ClusterShutdownState tracks the coordinated cluster shutdown progress in KV.
|
||||
type ClusterShutdownState struct {
|
||||
Initiator string `json:"initiator"`
|
||||
Phase string `json:"phase"`
|
||||
Started string `json:"started"`
|
||||
Timeout string `json:"timeout"`
|
||||
Force bool `json:"force"`
|
||||
NodesTotal int `json:"nodes_total"`
|
||||
NodesAcked map[string]string `json:"nodes_acked"`
|
||||
}
|
||||
|
||||
// WriteClusterShutdown writes the cluster shutdown state to KV.
|
||||
func (m *JetStreamManager) WriteClusterShutdown(state *ClusterShutdownState) error {
|
||||
if m.clusterKV == nil {
|
||||
return errors.New("cluster state KV not initialized")
|
||||
}
|
||||
data, err := json.Marshal(state)
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
_, err = m.clusterKV.Put("cluster.shutdown", data)
|
||||
return err
|
||||
}
|
||||
|
||||
// ReadClusterShutdown reads the cluster shutdown state from KV.
|
||||
func (m *JetStreamManager) ReadClusterShutdown() (*ClusterShutdownState, error) {
|
||||
if m.clusterKV == nil {
|
||||
return nil, errors.New("cluster state KV not initialized")
|
||||
}
|
||||
entry, err := m.clusterKV.Get("cluster.shutdown")
|
||||
if err != nil {
|
||||
return nil, err
|
||||
}
|
||||
var state ClusterShutdownState
|
||||
if err := json.Unmarshal(entry.Value(), &state); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
return &state, nil
|
||||
}
|
||||
|
||||
// DeleteClusterShutdown removes the cluster shutdown state from KV.
|
||||
func (m *JetStreamManager) DeleteClusterShutdown() error {
|
||||
if m.clusterKV == nil {
|
||||
return errors.New("cluster state KV not initialized")
|
||||
}
|
||||
err := m.clusterKV.Delete("cluster.shutdown")
|
||||
if err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
|
||||
return err
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
// WriteShutdownMarker writes a shutdown marker for the given node to the cluster-state KV.
|
||||
func (m *JetStreamManager) WriteShutdownMarker(nodeID string) error {
|
||||
if m.clusterKV == nil {
|
||||
|
||||
@@ -0,0 +1,276 @@
|
||||
package daemon
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"log/slog"
|
||||
"os"
|
||||
"os/exec"
|
||||
"time"
|
||||
|
||||
"github.com/mulgadc/hive/hive/utils"
|
||||
"github.com/nats-io/nats.go"
|
||||
)
|
||||
|
||||
// ShutdownRequest is sent by the coordinator to each daemon per phase.
|
||||
type ShutdownRequest struct {
|
||||
Phase string `json:"phase"`
|
||||
Force bool `json:"force"`
|
||||
Timeout int `json:"timeout_seconds"`
|
||||
}
|
||||
|
||||
// ShutdownACK is the response from a daemon after completing a shutdown phase.
|
||||
type ShutdownACK struct {
|
||||
Node string `json:"node"`
|
||||
Phase string `json:"phase"`
|
||||
Stopped []string `json:"stopped,omitempty"`
|
||||
Error string `json:"error,omitempty"`
|
||||
}
|
||||
|
||||
// ShutdownProgress is published by daemons during the DRAIN phase to report VM shutdown progress.
|
||||
type ShutdownProgress struct {
|
||||
Node string `json:"node"`
|
||||
Phase string `json:"phase"`
|
||||
Total int `json:"total"`
|
||||
Remaining int `json:"remaining"`
|
||||
}
|
||||
|
||||
// handleShutdownGate stops the API gateway and UI, then sets shuttingDown flag
|
||||
// so the daemon rejects new work. Phase: GATE.
|
||||
func (d *Daemon) handleShutdownGate(msg *nats.Msg) {
|
||||
var req ShutdownRequest
|
||||
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
||||
slog.Error("handleShutdownGate: failed to unmarshal request", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
slog.Info("Shutdown GATE phase starting", "node", d.node, "force", req.Force)
|
||||
|
||||
var stopped []string
|
||||
|
||||
// Stop AWSGW
|
||||
if d.config.HasService("awsgw") {
|
||||
if err := utils.StopProcess("awsgw"); err != nil {
|
||||
slog.Warn("Failed to stop awsgw", "error", err)
|
||||
} else {
|
||||
stopped = append(stopped, "awsgw")
|
||||
}
|
||||
}
|
||||
|
||||
// Stop UI
|
||||
if d.config.HasService("ui") {
|
||||
if err := utils.StopProcess("hive-ui"); err != nil {
|
||||
slog.Warn("Failed to stop hive-ui", "error", err)
|
||||
} else {
|
||||
stopped = append(stopped, "hive-ui")
|
||||
}
|
||||
}
|
||||
|
||||
// Set shuttingDown flag so daemon rejects new work
|
||||
d.shuttingDown.Store(true)
|
||||
|
||||
ack := ShutdownACK{
|
||||
Node: d.node,
|
||||
Phase: "gate",
|
||||
Stopped: stopped,
|
||||
}
|
||||
d.respondShutdownACK(msg, ack)
|
||||
slog.Info("Shutdown GATE phase complete", "node", d.node, "stopped", stopped)
|
||||
}
|
||||
|
||||
// handleShutdownDrain gracefully stops all running VMs, writes shutdown marker
|
||||
// and persists state. Phase: DRAIN.
|
||||
func (d *Daemon) handleShutdownDrain(msg *nats.Msg) {
|
||||
var req ShutdownRequest
|
||||
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
||||
slog.Error("handleShutdownDrain: failed to unmarshal request", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
slog.Info("Shutdown DRAIN phase starting", "node", d.node)
|
||||
|
||||
// Count total VMs
|
||||
d.Instances.Mu.Lock()
|
||||
total := len(d.Instances.VMS)
|
||||
d.Instances.Mu.Unlock()
|
||||
|
||||
// Publish initial progress
|
||||
d.publishShutdownProgress("drain", total, total)
|
||||
|
||||
// Stop all instances (graceful shutdown, no volume deletion)
|
||||
if total > 0 {
|
||||
d.Instances.Mu.Lock()
|
||||
vms := d.Instances.VMS
|
||||
d.Instances.Mu.Unlock()
|
||||
|
||||
if err := d.stopInstance(vms, false); err != nil {
|
||||
slog.Error("Failed to stop instances during DRAIN", "error", err)
|
||||
ack := ShutdownACK{
|
||||
Node: d.node,
|
||||
Phase: "drain",
|
||||
Error: err.Error(),
|
||||
}
|
||||
d.respondShutdownACK(msg, ack)
|
||||
return
|
||||
}
|
||||
}
|
||||
|
||||
// Publish final progress
|
||||
d.publishShutdownProgress("drain", total, 0)
|
||||
|
||||
// Write shutdown marker and persist state
|
||||
if d.jsManager != nil {
|
||||
if err := d.jsManager.WriteShutdownMarker(d.node); err != nil {
|
||||
slog.Error("Failed to write shutdown marker during DRAIN", "error", err)
|
||||
}
|
||||
}
|
||||
if err := d.WriteState(); err != nil {
|
||||
slog.Error("Failed to write state during DRAIN", "error", err)
|
||||
}
|
||||
|
||||
ack := ShutdownACK{
|
||||
Node: d.node,
|
||||
Phase: "drain",
|
||||
}
|
||||
d.respondShutdownACK(msg, ack)
|
||||
slog.Info("Shutdown DRAIN phase complete", "node", d.node, "vms_stopped", total)
|
||||
}
|
||||
|
||||
// handleShutdownStorage stops viperblock and cleans up orphan nbdkit processes. Phase: STORAGE.
|
||||
func (d *Daemon) handleShutdownStorage(msg *nats.Msg) {
|
||||
var req ShutdownRequest
|
||||
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
||||
slog.Error("handleShutdownStorage: failed to unmarshal request", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
slog.Info("Shutdown STORAGE phase starting", "node", d.node)
|
||||
|
||||
var stopped []string
|
||||
|
||||
if d.config.HasService("viperblock") {
|
||||
if err := utils.StopProcess("viperblock"); err != nil {
|
||||
slog.Warn("Failed to stop viperblock", "error", err)
|
||||
} else {
|
||||
stopped = append(stopped, "viperblock")
|
||||
}
|
||||
}
|
||||
|
||||
// Best-effort cleanup of orphaned nbdkit processes
|
||||
if err := exec.Command("pkill", "-f", "nbdkit").Run(); err != nil {
|
||||
slog.Debug("pkill nbdkit (best-effort)", "result", err)
|
||||
}
|
||||
|
||||
ack := ShutdownACK{
|
||||
Node: d.node,
|
||||
Phase: "storage",
|
||||
Stopped: stopped,
|
||||
}
|
||||
d.respondShutdownACK(msg, ack)
|
||||
slog.Info("Shutdown STORAGE phase complete", "node", d.node, "stopped", stopped)
|
||||
}
|
||||
|
||||
// handleShutdownPersist stops predastore. Phase: PERSIST.
|
||||
func (d *Daemon) handleShutdownPersist(msg *nats.Msg) {
|
||||
var req ShutdownRequest
|
||||
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
||||
slog.Error("handleShutdownPersist: failed to unmarshal request", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
slog.Info("Shutdown PERSIST phase starting", "node", d.node)
|
||||
|
||||
var stopped []string
|
||||
|
||||
if d.config.HasService("predastore") {
|
||||
if err := utils.StopProcess("predastore"); err != nil {
|
||||
slog.Warn("Failed to stop predastore", "error", err)
|
||||
} else {
|
||||
stopped = append(stopped, "predastore")
|
||||
}
|
||||
}
|
||||
|
||||
ack := ShutdownACK{
|
||||
Node: d.node,
|
||||
Phase: "persist",
|
||||
Stopped: stopped,
|
||||
}
|
||||
d.respondShutdownACK(msg, ack)
|
||||
slog.Info("Shutdown PERSIST phase complete", "node", d.node, "stopped", stopped)
|
||||
}
|
||||
|
||||
// handleShutdownInfra stops NATS and exits the daemon. Phase: INFRA.
|
||||
// No ACK is sent because NATS is going down — this is fire-and-forget from the coordinator.
|
||||
func (d *Daemon) handleShutdownInfra(msg *nats.Msg) {
|
||||
var req ShutdownRequest
|
||||
if err := json.Unmarshal(msg.Data, &req); err != nil {
|
||||
slog.Error("handleShutdownInfra: failed to unmarshal request", "error", err)
|
||||
return
|
||||
}
|
||||
|
||||
slog.Info("Shutdown INFRA phase starting", "node", d.node)
|
||||
|
||||
// Cancel context to stop heartbeat and other goroutines
|
||||
d.cancel()
|
||||
|
||||
// Unsubscribe from all NATS topics
|
||||
for topic, sub := range d.natsSubscriptions {
|
||||
slog.Debug("Unsubscribing from NATS", "topic", topic)
|
||||
if err := sub.Unsubscribe(); err != nil {
|
||||
slog.Warn("Failed to unsubscribe", "topic", topic, "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Shutdown cluster manager
|
||||
if d.clusterApp != nil {
|
||||
if err := d.clusterApp.Shutdown(); err != nil {
|
||||
slog.Warn("Failed to shutdown cluster manager", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// Close NATS connection
|
||||
d.natsConn.Close()
|
||||
|
||||
// Stop NATS server if this node runs it
|
||||
if d.config.HasService("nats") {
|
||||
if err := utils.StopProcess("nats"); err != nil {
|
||||
slog.Warn("Failed to stop nats", "error", err)
|
||||
} else {
|
||||
slog.Info("NATS server stopped", "node", d.node)
|
||||
}
|
||||
}
|
||||
|
||||
slog.Info("Shutdown INFRA phase complete, exiting", "node", d.node)
|
||||
|
||||
// Give a moment for log flush
|
||||
time.Sleep(500 * time.Millisecond)
|
||||
os.Exit(0)
|
||||
}
|
||||
|
||||
// respondShutdownACK marshals and sends a ShutdownACK response.
|
||||
func (d *Daemon) respondShutdownACK(msg *nats.Msg, ack ShutdownACK) {
|
||||
data, err := json.Marshal(ack)
|
||||
if err != nil {
|
||||
slog.Error("Failed to marshal shutdown ACK", "error", err)
|
||||
return
|
||||
}
|
||||
if err := msg.Respond(data); err != nil {
|
||||
slog.Error("Failed to respond with shutdown ACK", "error", err)
|
||||
}
|
||||
}
|
||||
|
||||
// publishShutdownProgress publishes a progress update during the DRAIN phase.
|
||||
func (d *Daemon) publishShutdownProgress(phase string, total, remaining int) {
|
||||
progress := ShutdownProgress{
|
||||
Node: d.node,
|
||||
Phase: phase,
|
||||
Total: total,
|
||||
Remaining: remaining,
|
||||
}
|
||||
data, err := json.Marshal(progress)
|
||||
if err != nil {
|
||||
return
|
||||
}
|
||||
if err := d.natsConn.Publish("hive.cluster.shutdown.progress", data); err != nil {
|
||||
slog.Debug("Failed to publish shutdown progress", "error", err)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,147 @@
|
||||
package daemon
|
||||
|
||||
import (
|
||||
"encoding/json"
|
||||
"testing"
|
||||
|
||||
"github.com/nats-io/nats.go"
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
// TestShutdownACKMarshal verifies JSON round-trip for ShutdownACK.
|
||||
func TestShutdownACKMarshal(t *testing.T) {
|
||||
ack := ShutdownACK{
|
||||
Node: "node1",
|
||||
Phase: "gate",
|
||||
Stopped: []string{"awsgw", "hive-ui"},
|
||||
}
|
||||
|
||||
data, err := json.Marshal(ack)
|
||||
require.NoError(t, err)
|
||||
|
||||
var decoded ShutdownACK
|
||||
err = json.Unmarshal(data, &decoded)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, ack.Node, decoded.Node)
|
||||
assert.Equal(t, ack.Phase, decoded.Phase)
|
||||
assert.Equal(t, ack.Stopped, decoded.Stopped)
|
||||
assert.Empty(t, decoded.Error)
|
||||
}
|
||||
|
||||
// TestShutdownACKWithError verifies JSON round-trip for ShutdownACK with an error.
|
||||
func TestShutdownACKWithError(t *testing.T) {
|
||||
ack := ShutdownACK{
|
||||
Node: "node2",
|
||||
Phase: "drain",
|
||||
Error: "failed to stop VMs",
|
||||
}
|
||||
|
||||
data, err := json.Marshal(ack)
|
||||
require.NoError(t, err)
|
||||
|
||||
var decoded ShutdownACK
|
||||
err = json.Unmarshal(data, &decoded)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "failed to stop VMs", decoded.Error)
|
||||
assert.Nil(t, decoded.Stopped)
|
||||
}
|
||||
|
||||
// TestShutdownRequestMarshal verifies JSON round-trip for ShutdownRequest.
|
||||
func TestShutdownRequestMarshal(t *testing.T) {
|
||||
req := ShutdownRequest{
|
||||
Phase: "drain",
|
||||
Force: true,
|
||||
Timeout: 120,
|
||||
}
|
||||
|
||||
data, err := json.Marshal(req)
|
||||
require.NoError(t, err)
|
||||
|
||||
var decoded ShutdownRequest
|
||||
err = json.Unmarshal(data, &decoded)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "drain", decoded.Phase)
|
||||
assert.True(t, decoded.Force)
|
||||
assert.Equal(t, 120, decoded.Timeout)
|
||||
}
|
||||
|
||||
// TestShutdownProgressMarshal verifies JSON round-trip for ShutdownProgress.
|
||||
func TestShutdownProgressMarshal(t *testing.T) {
|
||||
progress := ShutdownProgress{
|
||||
Node: "node1",
|
||||
Phase: "drain",
|
||||
Total: 5,
|
||||
Remaining: 3,
|
||||
}
|
||||
|
||||
data, err := json.Marshal(progress)
|
||||
require.NoError(t, err)
|
||||
|
||||
var decoded ShutdownProgress
|
||||
err = json.Unmarshal(data, &decoded)
|
||||
require.NoError(t, err)
|
||||
|
||||
assert.Equal(t, "node1", decoded.Node)
|
||||
assert.Equal(t, "drain", decoded.Phase)
|
||||
assert.Equal(t, 5, decoded.Total)
|
||||
assert.Equal(t, 3, decoded.Remaining)
|
||||
}
|
||||
|
||||
// TestClusterShutdownStateKVRoundTrip verifies cluster shutdown state can be stored and retrieved from KV.
|
||||
func TestClusterShutdownStateKVRoundTrip(t *testing.T) {
|
||||
nc, err := nats.Connect(sharedJSNATSURL)
|
||||
require.NoError(t, err)
|
||||
defer nc.Close()
|
||||
|
||||
jsm, err := NewJetStreamManager(nc, 1)
|
||||
require.NoError(t, err)
|
||||
err = jsm.InitClusterStateBucket()
|
||||
require.NoError(t, err)
|
||||
|
||||
state := &ClusterShutdownState{
|
||||
Initiator: "node1",
|
||||
Phase: "drain",
|
||||
Started: "2025-01-01T00:00:00Z",
|
||||
Timeout: "2m0s",
|
||||
Force: false,
|
||||
NodesTotal: 3,
|
||||
NodesAcked: map[string]string{
|
||||
"node1": "gate",
|
||||
"node2": "gate",
|
||||
},
|
||||
}
|
||||
|
||||
err = jsm.WriteClusterShutdown(state)
|
||||
require.NoError(t, err)
|
||||
|
||||
loaded, err := jsm.ReadClusterShutdown()
|
||||
require.NoError(t, err)
|
||||
require.NotNil(t, loaded)
|
||||
|
||||
assert.Equal(t, state.Initiator, loaded.Initiator)
|
||||
assert.Equal(t, state.Phase, loaded.Phase)
|
||||
assert.Equal(t, state.Started, loaded.Started)
|
||||
assert.Equal(t, state.NodesTotal, loaded.NodesTotal)
|
||||
assert.Equal(t, state.Force, loaded.Force)
|
||||
assert.Len(t, loaded.NodesAcked, 2)
|
||||
|
||||
// Cleanup
|
||||
err = jsm.DeleteClusterShutdown()
|
||||
require.NoError(t, err)
|
||||
}
|
||||
|
||||
// TestShuttingDownFlagSkipsVMStop verifies that the shuttingDown flag is respected.
|
||||
func TestShuttingDownFlagSkipsVMStop(t *testing.T) {
|
||||
d := &Daemon{}
|
||||
|
||||
// Default should be false
|
||||
assert.False(t, d.shuttingDown.Load())
|
||||
|
||||
// Set to true (as GATE phase does)
|
||||
d.shuttingDown.Store(true)
|
||||
assert.True(t, d.shuttingDown.Load())
|
||||
}
|
||||
@@ -50,6 +50,17 @@ has_service() {
|
||||
}
|
||||
echo "Services: $SERVICES"
|
||||
|
||||
# Detect multi-node cluster from config
|
||||
is_multinode() {
|
||||
local config_file="$CONFIG_DIR/hive.toml"
|
||||
if [ -f "$config_file" ]; then
|
||||
local node_count=$(grep -c '^\[nodes\.' "$config_file")
|
||||
[ "$node_count" -gt 1 ]
|
||||
else
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# Confirm configuration directory exists
|
||||
if [ ! -d "$CONFIG_DIR" ]; then
|
||||
echo "⚠️ Configuration directory $CONFIG_DIR does not exist."
|
||||
@@ -324,6 +335,36 @@ echo " export AWS_PROFILE=hive"
|
||||
echo " aws ec2 describe-instances"
|
||||
echo ""
|
||||
|
||||
# For multi-node clusters, check peer daemon health (best-effort)
|
||||
if is_multinode; then
|
||||
echo "Checking cluster peer health..."
|
||||
# Extract peer daemon hosts from config (host = "ip:port" under [nodes.X.daemon])
|
||||
peer_hosts=$(awk '
|
||||
/^\[nodes\./ { node=1; daemon=0 }
|
||||
node && /^\[.*\.daemon\]/ { daemon=1 }
|
||||
daemon && /^host/ { gsub(/[" ]/, "", $3); print $3; daemon=0 }
|
||||
' "$CONFIG_DIR/hive.toml")
|
||||
|
||||
for peer in $peer_hosts; do
|
||||
attempts=0
|
||||
max_attempts=5
|
||||
while [ $attempts -lt $max_attempts ]; do
|
||||
if curl -s --connect-timeout 3 "http://$peer/health" > /dev/null 2>&1; then
|
||||
echo " $peer: healthy"
|
||||
break
|
||||
fi
|
||||
((attempts++))
|
||||
if [ $attempts -lt $max_attempts ]; then
|
||||
sleep 2
|
||||
fi
|
||||
done
|
||||
if [ $attempts -ge $max_attempts ]; then
|
||||
echo " $peer: not responding (may still be starting)"
|
||||
fi
|
||||
done
|
||||
echo ""
|
||||
fi
|
||||
|
||||
# This will only be reached if air/daemon exits normally
|
||||
#echo ""
|
||||
#echo "🛑 Hive development environment stopped"
|
||||
|
||||
@@ -52,6 +52,26 @@ has_service() {
|
||||
echo "$SERVICES" | grep -qw "$svc"
|
||||
}
|
||||
|
||||
# Detect multi-node cluster from config
|
||||
is_multinode() {
|
||||
local config_file="$CONFIG_DIR/hive.toml"
|
||||
if [ -f "$config_file" ]; then
|
||||
local node_count=$(grep -c '^\[nodes\.' "$config_file")
|
||||
[ "$node_count" -gt 1 ]
|
||||
else
|
||||
return 1
|
||||
fi
|
||||
}
|
||||
|
||||
# For multi-node clusters, delegate to coordinated shutdown via NATS.
|
||||
# Only when called without arguments (default path). When a data-dir is
|
||||
# explicitly provided (e.g., E2E per-node cleanup), use per-service stop.
|
||||
if is_multinode && [ -z "$HIVE_FORCE_LOCAL_STOP" ] && [ -z "$1" ]; then
|
||||
echo "Multi-node cluster detected. Using coordinated shutdown..."
|
||||
echo " (Set HIVE_FORCE_LOCAL_STOP=1 to force per-service stop)"
|
||||
exec $PROJECT_ROOT/bin/hive admin cluster shutdown
|
||||
fi
|
||||
|
||||
# Function to stop service by PID file
|
||||
stop_service() {
|
||||
local name="$1"
|
||||
|
||||
Reference in New Issue
Block a user