fix(eks,ebs): EKS multi-node create-cluster + EBS/instance lifecycle hardening (#487)

* feat(ec2,ebs): harden volume attach/detach + instance lifecycle state machine

Batch of related EBS/instance-lifecycle correctness fixes verified together.

- u94vs: add compare-and-swap (casUpdate/casPut via nats KV Update revision)
  to the multi-writer lifecycle KV sites (WriteTerminatedInstance,
  WriteStoppedInstance, WriteClusterShutdown); make the teardown-reaper
  read-modify-write atomic via a merge callback so retries never re-run
  cleaner side effects.
- yi2ed: add ClaimStoppedInstance, a revision-checked delete-as-claim on the
  stopped bucket performed before Allocate/Insert/Run, so a forwarded start
  and a local fallback cannot double-start one instance onto the same volume.
- siv-473: make AttachVolume idempotent — return attached success when the
  volume is already attached to the requesting instance; a conflicting
  requested device returns InvalidParameterValue.
- siv-462: reconnect now re-asserts in-use for all attached volumes (boot and
  non-boot data volumes), not just boot, so a daemon restart no longer leaves
  stateful-pod data volumes split-brain; EFI pflash is skipped.
- 9w2mv: wait for the QMP DEVICE_DELETED event (bounded) before blockdev-del
  instead of a blind sleep, so a slow guest unplug no longer exhausts retries.
- soyah: poll (50ms) instead of fixed 1s sleeps in KillProcess and the
  viperblockd nbdkit confirm-running path, cutting volume-lifecycle latency.
- siv-481: a cross-node start forward that times out now falls back to a local
  start; this is safe because the yi2ed claim is the single serialization point.

* fix(eks): create per-account and leader KV buckets with cluster replica count

The EKS per-account KV bucket and the leader lease bucket were created at
Replicas=1 and never upgraded, since they're created lazily after
upgradeJetStreamReplicas runs at startup. On multi-node clusters this leaves
their single-replica Raft group transiently leaderless during formation,
causing KV timeouts that surface as a deterministic 503 on ListClusters.

Thread clusterSize through GetOrCreateAccountBucket/InitLeaderBucket (backed
by a new utils.GetOrCreateKVBucketWithReplicas) so both buckets are created
at max(clusterSize,1) replicas, matching every other node-count-scaled
JetStream stream.
This commit is contained in:
Josh Siviour
2026-07-09 16:49:31 +10:00
committed by GitHub
parent 20288dbe5d
commit ec9f0e5f2d
51 changed files with 1939 additions and 202 deletions
+22 -15
View File
@@ -175,9 +175,15 @@ type Daemon struct {
// initJetStream succeeds.
stateStore vm.StateStore
// Delay after QMP device_del before blockdev-del (default 1s, 0 in tests)
// Delay after QMP device_del before blockdev-del (default 1s, 0 in tests).
// Only used as a fallback when deviceDeletedTimeout is 0.
detachDelay time.Duration
// deviceDeletedTimeout bounds how long DetachVolume waits for QEMU's
// DEVICE_DELETED event after device_del before falling back to the
// blockdev-del retry loop (default 15s, 0 disables the wait in tests).
deviceDeletedTimeout time.Duration
// NATS connect retry options (nil uses defaults: 5min max, 500ms initial delay)
natsRetryOpts []utils.RetryOption
@@ -701,20 +707,21 @@ func NewDaemon(cfg *config.ClusterConfig) (*Daemon, error) {
}
d := &Daemon{
node: cfg.Node,
clusterConfig: cfg,
config: &nodeCfg,
resourceMgr: rm,
gpuProbe: gpuProbe,
gpuManager: gpuMgr,
ctx: ctx,
cancel: cancel,
vmMgr: vm.NewManager(),
natsSubscriptions: make(map[string]*nats.Subscription),
startTime: time.Now(),
detachDelay: 1 * time.Second,
requireNATSTimeout: 30 * time.Second,
exitFunc: os.Exit,
node: cfg.Node,
clusterConfig: cfg,
config: &nodeCfg,
resourceMgr: rm,
gpuProbe: gpuProbe,
gpuManager: gpuMgr,
ctx: ctx,
cancel: cancel,
vmMgr: vm.NewManager(),
natsSubscriptions: make(map[string]*nats.Subscription),
startTime: time.Now(),
detachDelay: 1 * time.Second,
deviceDeletedTimeout: 15 * time.Second,
requireNATSTimeout: 30 * time.Second,
exitFunc: os.Exit,
}
// Initialise peersReachable true so the first probe tick never fires a
// spurious reconcileOnHeal at startup. Mode() still requires natsConnected
+12 -19
View File
@@ -25,11 +25,9 @@ import (
// startStoppedForwardTimeout bounds the wait for ec2.start.{nodeId} to respond.
// StartStoppedInstance does volume rehydrate + QEMU launch + QMP handshake,
// which can take 10-20s on a cold start. Match the awsgw upstream 30s budget
// so a slow-but-alive target node never trips the fallback path — falling back
// while the target is still launching causes a duplicate Run that collides on
// the deterministic tap name (TUNSETIFF: Device or resource busy).
const startStoppedForwardTimeout = 30 * time.Second
// which can take 10-20s on a cold start. Match the awsgw upstream 30s budget.
// A var (not const) so tests can shrink it instead of sleeping real seconds.
var startStoppedForwardTimeout = 30 * time.Second
// handleSetInstanceTags applies a create-tags/delete-tags mutation to a running
// instance: central store first, then the record under the manager lock, so a
@@ -477,14 +475,11 @@ func (d *Daemon) handleEC2DescribeInstanceStatus(msg *nats.Msg) {
// handleEC2StartStoppedInstance handles the generic ec2.start queue-group topic.
// It reads the stopped instance's LastNode from shared KV and forwards the
// request to ec2.start.{lastNode} when the instance last ran on a different
// node. This keeps instances on their original node so the per-node resource
// manager sees the correct allocation. Local fallback fires only when the
// targeted node has no active subscriber (ErrNoResponders — node down or not
// yet recovered) or returns InsufficientInstanceCapacity. A timeout from a
// reachable-but-slow target is surfaced as ServerInternal so the caller can
// retry; falling back in that case races the still-running launch on the
// target and collides on the deterministic tap name.
// request to ec2.start.{lastNode} to keep instances on their original node.
// Local fallback fires on any forward failure — no responders, capacity, or
// timeout — because StartStoppedInstance's ClaimStoppedInstance call is the
// single cluster-wide serialization point, so a losing racer bails out
// cleanly instead of double-starting.
func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) {
// Peek at the instance ID without full unmarshal — we only need it for the
// LastNode lookup. The full unmarshal happens inside StartStoppedInstance.
@@ -532,13 +527,11 @@ func (d *Daemon) handleEC2StartStoppedInstance(msg *nats.Msg) {
slog.Warn("ec2.start: original node has no subscriber, starting locally",
"instanceId", peek.InstanceID, "lastNode", lastNode, "err", err)
} else {
// Timeout or other transport error from a node whose subscription
// did exist. Do NOT fall back — the target may still be launching
// the VM, and a duplicate Run here would collide on tap setup.
slog.Error("ec2.start: forward to original node failed, not falling back",
// Timeout or other transport error — e.g. lastNode crashed without
// a clean NATS disconnect, so no ErrNoResponders fires. The atomic
// claim in StartStoppedInstance makes a local retry safe either way.
slog.Warn("ec2.start: forward to original node failed, attempting local start",
"instanceId", peek.InstanceID, "lastNode", lastNode, "err", err)
respondWithError(msg, awserrors.ErrorServerInternal)
return
}
}
@@ -52,6 +52,15 @@ func (m *memStoppedStore) DeleteStoppedInstance(id string) error {
return nil
}
func (m *memStoppedStore) ClaimStoppedInstance(id string) (*vm.VM, error) {
v, ok := m.instances[id]
if !ok {
return nil, vm.ErrStoppedInstanceClaimed
}
delete(m.instances, id)
return v, nil
}
func (m *memStoppedStore) WriteTerminatedInstance(string, *vm.VM) error { return nil }
// tagTestDaemon returns a test daemon with an in-memory central tag store, an
@@ -44,6 +44,8 @@ type fakeStateStore struct {
deleteStoppedFailFirst bool
deleteStoppedAttempts int
claimStoppedErr error
terminated map[string]*vm.VM
writeTerminatedErr error
listTerminatedErr error
@@ -104,6 +106,23 @@ func (f *fakeStateStore) DeleteStoppedInstance(id string) error {
return nil
}
// ClaimStoppedInstance mimics the real atomic-delete claim for tests: under
// the store lock, remove and return the entry, or claimStoppedErr /
// vm.ErrStoppedInstanceClaimed if it is already gone or a test forced an error.
func (f *fakeStateStore) ClaimStoppedInstance(id string) (*vm.VM, error) {
if f.claimStoppedErr != nil {
return nil, f.claimStoppedErr
}
f.mu.Lock()
defer f.mu.Unlock()
v, ok := f.stopped[id]
if !ok {
return nil, vm.ErrStoppedInstanceClaimed
}
delete(f.stopped, id)
return v, nil
}
func (f *fakeStateStore) ListStoppedInstances() ([]*vm.VM, error) {
if f.listStoppedErr != nil {
return nil, f.listStoppedErr
@@ -147,6 +166,19 @@ func (f *fakeStateStore) DeleteTerminatedInstance(id string) error {
return nil
}
// UpdateTerminatedInstance mimics the real CAS semantics for tests: mutate
// runs under the store lock against the stored value.
func (f *fakeStateStore) UpdateTerminatedInstance(id string, mutate func(*vm.VM)) (*vm.VM, error) {
f.mu.Lock()
defer f.mu.Unlock()
v, ok := f.terminated[id]
if !ok {
return nil, errors.New("terminated instance not found")
}
mutate(v)
return v, nil
}
var _ vm.StateStore = (*fakeStateStore)(nil)
// daemonWithFakeStateStore returns a daemon wired with an in-memory NATS
@@ -257,6 +289,81 @@ func TestHandleEC2StartStoppedInstance_InstanceTypeUnknown(t *testing.T) {
assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, decodeError(t, reply.Data)["Code"])
}
// withShortForwardTimeout shrinks startStoppedForwardTimeout for the duration
// of a test so a forced forward timeout doesn't cost real wall-clock seconds.
func withShortForwardTimeout(t *testing.T, d time.Duration) {
t.Helper()
orig := startStoppedForwardTimeout
startStoppedForwardTimeout = d
t.Cleanup(func() { startStoppedForwardTimeout = orig })
}
// TestHandleEC2StartStoppedInstance_ForwardTimeoutFallsBackLocally pins
// siv-481: a forward to LastNode that times out (as opposed to an immediate
// ErrNoResponders) must still fall back to a local start attempt instead of
// surfacing a bare ServerInternal. The target subscriber below is alive but
// silent, so nats: timeout is the only error the forward can produce — proof
// that the fallback path, not ErrNoResponders handling, is what fires here.
// An unresolvable instance type turns "local start was attempted" into a
// distinct, assertable response code (InsufficientInstanceCapacity) instead
// of colliding with the old no-fallback ServerInternal response.
func TestHandleEC2StartStoppedInstance_ForwardTimeoutFallsBackLocally(t *testing.T) {
withShortForwardTimeout(t, 50*time.Millisecond)
store := newFakeStateStore()
v := stoppedVMFixture("i-timeout-fallback", testAccountID)
v.InstanceType = "definitely.not.a.real.type"
v.LastNode = "node-other"
store.stopped[v.ID] = v
d := daemonWithFakeStateStore(t, store)
// Simulate a live-but-unresponsive original node: subscribed, so no
// ErrNoResponders, but it never replies, so the forward times out.
silentSub, err := d.natsConn.Subscribe("ec2.start.node-other", func(*nats.Msg) {})
require.NoError(t, err)
t.Cleanup(func() { _ = silentSub.Unsubscribe() })
body, _ := json.Marshal(handlers_ec2_instance.StartStoppedInstanceInput{InstanceID: v.ID})
reply := requestHandler(t, d.natsConn, "ec2.start.test5", d.handleEC2StartStoppedInstance, testAccountID, body)
assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, decodeError(t, reply.Data)["Code"],
"a forward timeout must fall back to a local start attempt, not a bare ServerInternal")
}
// TestHandleEC2StartStoppedInstance_ForwardTimeoutAfterRemoteClaim_NoDoubleStart
// pins the other half of siv-481: if the forward times out on the caller's
// side AFTER the original node already won the atomic claim (removed the
// record from shared KV) and kept working, the caller's local fallback must
// not double-start the instance. It should observe the record already gone
// and fail cleanly, and must never insert a second copy into its own vmMgr.
func TestHandleEC2StartStoppedInstance_ForwardTimeoutAfterRemoteClaim_NoDoubleStart(t *testing.T) {
withShortForwardTimeout(t, 50*time.Millisecond)
store := newFakeStateStore()
v := stoppedVMFixture("i-race-claim", testAccountID)
v.LastNode = "node-other"
store.stopped[v.ID] = v
d := daemonWithFakeStateStore(t, store)
// Simulate the original node winning the claim (atomically removing the
// shared-KV record, as StartStoppedInstance's ClaimStoppedInstance does)
// but never replying — e.g. still mid-launch when the caller's forward
// budget expires.
claimingSub, err := d.natsConn.Subscribe("ec2.start.node-other", func(*nats.Msg) {
_, claimErr := store.ClaimStoppedInstance(v.ID)
assert.NoError(t, claimErr)
})
require.NoError(t, err)
t.Cleanup(func() { _ = claimingSub.Unsubscribe() })
body, _ := json.Marshal(handlers_ec2_instance.StartStoppedInstanceInput{InstanceID: v.ID})
reply := requestHandler(t, d.natsConn, "ec2.start.test6", d.handleEC2StartStoppedInstance, testAccountID, body)
assert.Equal(t, awserrors.ErrorInvalidInstanceIDNotFound, decodeError(t, reply.Data)["Code"],
"local fallback must see the record already claimed and fail without double-starting")
_, found := d.vmMgr.Get(v.ID)
assert.False(t, found, "local fallback must not insert a second running copy of an already-claimed instance")
}
// --- handleEC2TerminateStoppedInstance ---
func TestHandleEC2TerminateStoppedInstance_LoadError(t *testing.T) {
+20
View File
@@ -59,6 +59,26 @@ func (d *Daemon) handleAttachVolume(ctx context.Context, msg *nats.Msg, command
}
if volCfg.VolumeMetadata.State != "available" {
if volCfg.VolumeMetadata.AttachedInstance == command.ID {
if command.AttachVolumeData.Device != "" &&
command.AttachVolumeData.Device != volCfg.VolumeMetadata.DeviceName {
slog.ErrorContext(ctx, "AttachVolume: requested device conflicts with existing attachment",
"volumeId", volumeID, "instanceId", command.ID,
"requestedDevice", command.AttachVolumeData.Device,
"attachedDevice", volCfg.VolumeMetadata.DeviceName)
respondWithError(msg, attachDetachErrorCode(vm.ErrVolumeDeviceMismatch))
return
}
// Volume is already attached to this instance (e.g. a CSI
// ControllerPublishVolume retry after a slow first attach).
// Treat as an idempotent success instead of VolumeInUse.
slog.InfoContext(ctx, "AttachVolume: volume already attached to requesting instance, returning idempotent success",
"volumeId", volumeID, "instanceId", command.ID, "device", volCfg.VolumeMetadata.DeviceName)
d.respondWithVolumeAttachment(msg, volumeID, command.ID, volCfg.VolumeMetadata.DeviceName, "attached")
return
}
slog.ErrorContext(ctx, "AttachVolume: volume not available",
"volumeId", volumeID, "state", volCfg.VolumeMetadata.State)
respondWithError(msg, awserrors.ErrorVolumeInUse)
@@ -1,13 +1,25 @@
package daemon
import (
"encoding/json"
"errors"
"fmt"
"strings"
"testing"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
awss3 "github.com/aws/aws-sdk-go/service/s3"
"github.com/mulgadc/spinifex/spinifex/awserrors"
"github.com/mulgadc/spinifex/spinifex/objectstore"
"github.com/mulgadc/spinifex/spinifex/qmp"
"github.com/mulgadc/spinifex/spinifex/types"
"github.com/mulgadc/spinifex/spinifex/vm"
"github.com/mulgadc/viperblock/viperblock"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// TestAttachDetachErrorCode locks down the manager-error → AWS-API-code
@@ -86,3 +98,219 @@ func TestAttachDetachErrorCode(t *testing.T) {
})
}
}
// seedVolumeConfig writes a minimal VolumeConfig to config.json, matching
// the seeding pattern used by the handleAttachVolume tests in
// daemon_handlers_test.go.
func seedVolumeConfig(t *testing.T, store *objectstore.MemoryObjectStore, volumeID string, meta viperblock.VolumeMetadata) {
t.Helper()
wrapper := struct {
VolumeConfig viperblock.VolumeConfig `json:"VolumeConfig"`
}{
VolumeConfig: viperblock.VolumeConfig{VolumeMetadata: meta},
}
data, err := json.Marshal(wrapper)
require.NoError(t, err)
_, err = store.PutObject(t.Context(), &awss3.PutObjectInput{
Bucket: aws.String("test-bucket"),
Key: aws.String(volumeID + "/config.json"),
Body: strings.NewReader(string(data)),
})
require.NoError(t, err)
}
// TestAttachVolume_IdempotentSameInstance verifies that re-attaching a
// volume already attached to the requesting instance (e.g. a CSI
// ControllerPublishVolume retry after a slow first attach) short-circuits
// to an idempotent "attached" success instead of VolumeInUse, and does not
// invoke vm.Manager.AttachVolume (the fake QMPClient on the seeded instance
// has no live socket, so a real AttachVolume call would fail rather than
// silently succeed with the pre-existing device).
func TestAttachVolume_IdempotentSameInstance(t *testing.T) {
tests := []struct {
name string
requestedDevice string
}{
{name: "no device specified", requestedDevice: ""},
{name: "device matches existing attachment", requestedDevice: "/dev/sdf"},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
daemon, store := createFullTestDaemonWithStore(t, sharedNATSURL)
instanceID := "i-attach-idempotent-" + strings.ReplaceAll(tt.name, " ", "-")
volumeID := "vol-idempotent-" + strings.ReplaceAll(tt.name, " ", "-")
instance := &vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
InstanceType: getTestInstanceType(t),
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.vmMgr.Insert(instance)
seedVolumeConfig(t, store, volumeID, viperblock.VolumeMetadata{
VolumeID: volumeID,
SizeGiB: 10,
State: "in-use",
TenantID: testAccountID,
AttachedInstance: instanceID,
DeviceName: "/dev/sdf",
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
daemon.handleEC2Events,
)
require.NoError(t, err)
defer sub.Unsubscribe()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
AttachVolume: true,
},
AttachVolumeData: &types.AttachVolumeData{
VolumeID: volumeID,
Device: tt.requestedDevice,
},
}
cmdData, err := json.Marshal(command)
require.NoError(t, err)
resp, err := natsRequest(daemon.natsConn,
fmt.Sprintf("ec2.cmd.%s", instanceID),
cmdData,
5*time.Second,
)
require.NoError(t, err)
var attachment ec2.VolumeAttachment
require.NoError(t, json.Unmarshal(resp.Data, &attachment))
assert.Equal(t, volumeID, *attachment.VolumeId)
assert.Equal(t, instanceID, *attachment.InstanceId)
assert.Equal(t, "/dev/sdf", *attachment.Device)
assert.Equal(t, "attached", *attachment.State)
})
}
}
// TestAttachVolume_InUseDifferentInstance is a regression test locking down
// that a volume attached to a DIFFERENT instance than the requester still
// returns VolumeInUse unchanged, distinguishing it from the same-instance
// idempotent short-circuit.
func TestAttachVolume_InUseDifferentInstance(t *testing.T) {
daemon, store := createFullTestDaemonWithStore(t, sharedNATSURL)
instanceID := "i-attach-vol-other-instance"
otherInstanceID := "i-already-attached-elsewhere"
volumeID := "vol-in-use-other-instance"
instance := &vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
InstanceType: getTestInstanceType(t),
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.vmMgr.Insert(instance)
seedVolumeConfig(t, store, volumeID, viperblock.VolumeMetadata{
VolumeID: volumeID,
SizeGiB: 10,
State: "in-use",
TenantID: testAccountID,
AttachedInstance: otherInstanceID,
DeviceName: "/dev/sdf",
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
daemon.handleEC2Events,
)
require.NoError(t, err)
defer sub.Unsubscribe()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
AttachVolume: true,
},
AttachVolumeData: &types.AttachVolumeData{
VolumeID: volumeID,
},
}
cmdData, err := json.Marshal(command)
require.NoError(t, err)
resp, err := natsRequest(daemon.natsConn,
fmt.Sprintf("ec2.cmd.%s", instanceID),
cmdData,
5*time.Second,
)
require.NoError(t, err)
assert.Contains(t, string(resp.Data), "VolumeInUse")
}
// TestAttachVolume_IdempotentSameInstance_DeviceMismatch verifies that a
// same-instance re-attach requesting a DIFFERENT device than the one
// already attached is treated as a real CSI conflict (InvalidParameterValue
// via vm.ErrVolumeDeviceMismatch), not silently echoed back or accepted.
func TestAttachVolume_IdempotentSameInstance_DeviceMismatch(t *testing.T) {
daemon, store := createFullTestDaemonWithStore(t, sharedNATSURL)
instanceID := "i-attach-device-mismatch"
volumeID := "vol-device-mismatch"
instance := &vm.VM{
ID: instanceID,
Status: vm.StateRunning,
AccountID: testAccountID,
InstanceType: getTestInstanceType(t),
Instance: &ec2.Instance{},
QMPClient: &qmp.QMPClient{},
}
daemon.vmMgr.Insert(instance)
seedVolumeConfig(t, store, volumeID, viperblock.VolumeMetadata{
VolumeID: volumeID,
SizeGiB: 10,
State: "in-use",
TenantID: testAccountID,
AttachedInstance: instanceID,
DeviceName: "/dev/sdf",
})
sub, err := daemon.natsConn.Subscribe(
fmt.Sprintf("ec2.cmd.%s", instanceID),
daemon.handleEC2Events,
)
require.NoError(t, err)
defer sub.Unsubscribe()
command := types.EC2InstanceCommand{
ID: instanceID,
Attributes: types.EC2CommandAttributes{
AttachVolume: true,
},
AttachVolumeData: &types.AttachVolumeData{
VolumeID: volumeID,
Device: "/dev/sdg",
},
}
cmdData, err := json.Marshal(command)
require.NoError(t, err)
resp, err := natsRequest(daemon.natsConn,
fmt.Sprintf("ec2.cmd.%s", instanceID),
cmdData,
5*time.Second,
)
require.NoError(t, err)
assert.Contains(t, string(resp.Data), awserrors.ErrorInvalidParameterValue)
}
+7 -5
View File
@@ -119,7 +119,8 @@ func createTestDaemon(t *testing.T, natsURL string) *Daemon {
require.NoError(t, err, "Failed to connect to NATS")
daemon.natsConn = nc
daemon.detachDelay = 0 // Skip sleep in tests
daemon.detachDelay = 0 // Skip sleep in tests
daemon.deviceDeletedTimeout = 0 // Skip DEVICE_DELETED wait in tests
// Initialize services (needed for handler tests).
// jsManager is nil here; pass a nil literal to keep the StoppedInstanceStore
@@ -133,10 +134,11 @@ func createTestDaemon(t *testing.T, natsURL string) *Daemon {
// AttachVolume / DetachVolume manager methods enough plumbing to drive
// ebs.mount/unmount over NATS using the test's connection.
daemon.vmMgr.SetDeps(vm.Deps{
NodeID: daemon.node,
VolumeMounter: newVolumeMounterAdapter(daemon.natsConn, daemon.node, daemon.volumeService),
VolumeStateUpdater: daemon.volumeService,
DetachDelay: daemon.detachDelay,
NodeID: daemon.node,
VolumeMounter: newVolumeMounterAdapter(daemon.natsConn, daemon.node, daemon.volumeService),
VolumeStateUpdater: daemon.volumeService,
DetachDelay: daemon.detachDelay,
DeviceDeletedTimeout: daemon.deviceDeletedTimeout,
})
t.Cleanup(func() {
+3
View File
@@ -57,8 +57,10 @@ func (d *Daemon) buildEKSServiceDeps() handlers_eks.EKSServiceDeps {
}
internalSuffix := ""
clusterSize := 1
if d.clusterConfig != nil {
internalSuffix = d.clusterConfig.AWS.InternalSuffix
clusterSize = len(d.clusterConfig.Nodes)
}
gatewayCA := ""
@@ -78,6 +80,7 @@ func (d *Daemon) buildEKSServiceDeps() handlers_eks.EKSServiceDeps {
GatewayBaseURL: d.resolveGatewayBaseURL(),
Region: d.config.Region,
HolderID: d.node,
ClusterSize: clusterSize,
InternalSuffix: internalSuffix,
SystemGatewayURL: d.resolveSystemGatewayBaseURL(),
SystemAccessKey: d.config.Predastore.AccessKey,
+242 -13
View File
@@ -252,6 +252,146 @@ func (m *JetStreamManager) recoverTerminatedKVBucket() (nats.KeyValue, error) {
}, &m.terminatedKV, TerminatedInstanceBucketVersion)
}
// maxCASRetries bounds the optimistic-concurrency retry loop in casUpdate.
// A conflict this many times in a row means sustained write contention on
// the key, not a single overlapping update — callers should surface an error.
const maxCASRetries = 5
// isRevisionConflict reports whether err is the nats.go "wrong last
// sequence" error that KeyValue.Update/Create return when the key's
// revision no longer matches what the caller expected, i.e. a concurrent
// writer won the race. nats.go surfaces this as nats.ErrKeyExists for both
// Update (stale revision) and Create (key already present).
func isRevisionConflict(err error) bool {
return errors.Is(err, nats.ErrKeyExists)
}
// casUpdate performs an optimistic-concurrency read-modify-write against kv:
// Get the current entry (or start from a zero value when absent and
// createIfAbsent is set), apply mutate to the decoded value, then commit via
// Update/Create using the observed revision. A concurrent writer that changes
// the key between Get and Update/Create is detected via a revision conflict
// and retried, bounded by maxCASRetries.
func casUpdate[T any](kv nats.KeyValue, key string, mutate func(*T), createIfAbsent bool) (*T, error) {
var lastErr error
for range maxCASRetries {
var value T
var revision uint64
entry, err := kv.Get(key)
switch {
case err == nil:
if uerr := json.Unmarshal(entry.Value(), &value); uerr != nil {
return nil, uerr
}
revision = entry.Revision()
case errors.Is(err, nats.ErrKeyNotFound) && createIfAbsent:
revision = 0
default:
return nil, err
}
mutate(&value)
data, merr := json.Marshal(&value)
if merr != nil {
return nil, merr
}
if revision == 0 {
_, err = kv.Create(key, data)
} else {
_, err = kv.Update(key, data, revision)
}
if err == nil {
return &value, nil
}
if isRevisionConflict(err) {
lastErr = err
continue
}
return nil, err
}
return nil, fmt.Errorf("CAS update exhausted %d retries for key %s: %w", maxCASRetries, key, lastErr)
}
// casPut writes data to key under kv using optimistic concurrency: it reads
// the current revision (or treats the key as absent, using Create, when
// createIfAbsent is set) then commits via Update/Create, retrying on a
// revision conflict up to maxCASRetries. This is the "replace wholesale"
// counterpart to casUpdate, for values (like vm.VM, which embeds a
// sync.Mutex) that cannot be safely decoded into and copied out of a shared
// struct value.
func casPut(kv nats.KeyValue, key string, data []byte, createIfAbsent bool) error {
var lastErr error
for range maxCASRetries {
var revision uint64
entry, err := kv.Get(key)
switch {
case err == nil:
revision = entry.Revision()
case errors.Is(err, nats.ErrKeyNotFound) && createIfAbsent:
revision = 0
default:
return err
}
if revision == 0 {
_, err = kv.Create(key, data)
} else {
_, err = kv.Update(key, data, revision)
}
if err == nil {
return nil
}
if isRevisionConflict(err) {
lastErr = err
continue
}
return err
}
return fmt.Errorf("CAS put exhausted %d retries for key %s: %w", maxCASRetries, key, lastErr)
}
// casClaim atomically removes key from kv, decoding its current value into T
// first — the delete-side counterpart to casPut/casUpdate. It is the
// primitive an exclusive "claim" is built on: at most one caller can ever
// observe a successful delete for a given revision, so at most one caller
// gets back a non-nil value. A concurrent writer that changes the key
// between Get and Delete (another claim, or an unrelated update) is
// detected via a revision conflict and retried, bounded by maxCASRetries,
// so a losing racer only fails outright when the key is truly gone
// (notFound) or retries are exhausted.
func casClaim[T any](kv nats.KeyValue, key string) (value *T, notFound bool, err error) {
var lastErr error
for range maxCASRetries {
entry, gerr := kv.Get(key)
if gerr != nil {
if errors.Is(gerr, nats.ErrKeyNotFound) {
return nil, true, nil
}
return nil, false, gerr
}
var v T
if uerr := json.Unmarshal(entry.Value(), &v); uerr != nil {
return nil, false, uerr
}
if derr := kv.Delete(key, nats.LastRevision(entry.Revision())); derr != nil {
if isRevisionConflict(derr) {
lastErr = derr
continue
}
return nil, false, derr
}
return &v, false, nil
}
return nil, false, fmt.Errorf("CAS claim exhausted %d retries for key %s: %w", maxCASRetries, key, lastErr)
}
// Heartbeat represents a daemon's periodic health status published to cluster KV.
//
// AvailableVCPU / AvailableMem are observability-only (host - allocated,
@@ -313,19 +453,31 @@ type ClusterShutdownState struct {
NodesAcked map[string]string `json:"nodes_acked"`
}
// WriteClusterShutdown writes the cluster shutdown state to KV.
// WriteClusterShutdown writes the cluster shutdown state to KV, replacing
// any existing value. Uses CAS internally (bounded retry) so a write racing
// a concurrent update is detected and retried rather than silently applied
// out of order.
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)
_, err := casUpdate(m.clusterKV, "cluster.shutdown", func(s *ClusterShutdownState) {
*s = *state
}, true)
return err
}
// UpdateClusterShutdown atomically applies mutate to the current cluster
// shutdown state (e.g. merging a node's ack into NodesAcked) and writes it
// back with optimistic concurrency, retrying on a concurrent writer's
// revision conflict. Returns nats.ErrKeyNotFound if no shutdown is in progress.
func (m *JetStreamManager) UpdateClusterShutdown(mutate func(*ClusterShutdownState)) (*ClusterShutdownState, error) {
if m.clusterKV == nil {
return nil, errors.New("cluster state KV not initialized")
}
return casUpdate(m.clusterKV, "cluster.shutdown", mutate, false)
}
// ReadClusterShutdown reads the cluster shutdown state from KV.
func (m *JetStreamManager) ReadClusterShutdown() (*ClusterShutdownState, error) {
if m.clusterKV == nil {
@@ -628,7 +780,12 @@ func (m *JetStreamManager) UpdateReplicas(newReplicas int) error {
return nil
}
// WriteStoppedInstance writes a stopped instance to the shared KV store.
// WriteStoppedInstance writes a stopped instance to the shared KV store,
// replacing any existing record for instanceID. Uses CAS internally (bounded
// retry) so a write racing a concurrent update is detected and retried
// rather than silently overwriting it out of order. This is the substrate
// callers that read-modify-write a stopped instance (e.g. tag/attribute
// mutations) build on.
func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm.VM) error {
if m.kv == nil {
return errors.New("KV bucket not initialized")
@@ -640,7 +797,7 @@ func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm.
}
key := StoppedInstancePrefix + instanceID
_, err = m.kv.Put(key, jsonData)
err = casPut(m.kv, key, jsonData, true)
if err != nil {
if isStreamUnavailable(err) {
slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteStoppedInstance", "key", key, "err", err)
@@ -648,7 +805,7 @@ func (m *JetStreamManager) WriteStoppedInstance(instanceID string, instance *vm.
if recoverErr != nil {
return err
}
if _, retryErr := kv.Put(key, jsonData); retryErr != nil {
if retryErr := casPut(kv, key, jsonData, true); retryErr != nil {
return retryErr
}
slog.Debug("Wrote stopped instance to JetStream KV (after recovery)", "key", key, "instanceId", instanceID)
@@ -731,6 +888,45 @@ func (m *JetStreamManager) DeleteStoppedInstance(instanceID string) error {
return nil
}
// ClaimStoppedInstance atomically removes instanceID's record from the
// shared KV store and returns the VM it held, so that at most one caller
// across the cluster can ever win a race to (re)launch the same stopped
// instance. It is the first step of StartStoppedInstance, replacing the
// unguarded Load+Delete pair that previously let two callers both observe
// the instance as claimable. A losing racer (this node's local fallback
// racing a forwarded call, two nodes racing the same forwarded call, or a
// retry after the instance was already claimed) gets vm.ErrStoppedInstanceClaimed
// instead of a VM and must not proceed to allocate resources or launch qemu.
func (m *JetStreamManager) ClaimStoppedInstance(instanceID string) (*vm.VM, error) {
if m.kv == nil {
return nil, errors.New("KV bucket not initialized")
}
key := StoppedInstancePrefix + instanceID
instance, notFound, err := casClaim[vm.VM](m.kv, key)
if err != nil {
if isStreamUnavailable(err) {
slog.Warn("KV stream unavailable, attempting recovery", "operation", "ClaimStoppedInstance", "key", key, "err", err)
kv, recoverErr := m.recoverKVBucket()
if recoverErr != nil {
return nil, err
}
instance, notFound, err = casClaim[vm.VM](kv, key)
if err != nil {
return nil, err
}
} else {
return nil, err
}
}
if notFound {
return nil, vm.ErrStoppedInstanceClaimed
}
slog.Debug("Claimed stopped instance from JetStream KV", "key", key, "instanceId", instanceID)
return instance, nil
}
// ListStoppedInstances returns all stopped instances from the shared KV store.
func (m *JetStreamManager) ListStoppedInstances() ([]*vm.VM, error) {
if m.kv == nil {
@@ -791,8 +987,14 @@ func (m *JetStreamManager) ListStoppedInstances() ([]*vm.VM, error) {
return instances, nil
}
// WriteTerminatedInstance writes a terminated instance to the terminated KV bucket.
// The entry will auto-expire after the bucket's TTL (1 hour).
// WriteTerminatedInstance writes a terminated instance to the terminated KV
// bucket, replacing any existing record for instanceID. The entry will
// auto-expire after the bucket's TTL (1 hour). Uses CAS internally (bounded
// retry) so a write racing a concurrent update (e.g. the teardown reaper
// advancing the Teardown map) is detected and retried rather than silently
// overwriting it out of order. Callers that need to merge into the current
// record instead of replacing it wholesale should use
// UpdateTerminatedInstance.
func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance *vm.VM) error {
if m.terminatedKV == nil {
return errors.New("terminated instance KV bucket not initialized")
@@ -804,7 +1006,7 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance *
}
key := TerminatedInstancePrefix + instanceID
_, err = m.terminatedKV.Put(key, jsonData)
err = casPut(m.terminatedKV, key, jsonData, true)
if err != nil {
if isStreamUnavailable(err) {
slog.Warn("KV stream unavailable, attempting recovery", "operation", "WriteTerminatedInstance", "key", key, "err", err)
@@ -812,7 +1014,7 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance *
if recoverErr != nil {
return err
}
if _, retryErr := kv.Put(key, jsonData); retryErr != nil {
if retryErr := casPut(kv, key, jsonData, true); retryErr != nil {
return retryErr
}
slog.Debug("Wrote terminated instance to JetStream KV (after recovery)", "key", key, "instanceId", instanceID)
@@ -825,6 +1027,33 @@ func (m *JetStreamManager) WriteTerminatedInstance(instanceID string, instance *
return nil
}
// UpdateTerminatedInstance atomically applies mutate to the current
// KV-stored terminated record for instanceID and writes it back using
// optimistic concurrency (CAS), retrying on a concurrent writer's revision
// conflict. Used by the teardown reaper to merge per-dependent Teardown
// progress without clobbering marks written by a concurrent update to the
// same record. Returns nats.ErrKeyNotFound if no record exists yet.
func (m *JetStreamManager) UpdateTerminatedInstance(instanceID string, mutate func(*vm.VM)) (*vm.VM, error) {
if m.terminatedKV == nil {
return nil, errors.New("terminated instance KV bucket not initialized")
}
key := TerminatedInstancePrefix + instanceID
updated, err := casUpdate(m.terminatedKV, key, mutate, false)
if err != nil {
if isStreamUnavailable(err) {
slog.Warn("KV stream unavailable, attempting recovery", "operation", "UpdateTerminatedInstance", "key", key, "err", err)
kv, recoverErr := m.recoverTerminatedKVBucket()
if recoverErr != nil {
return nil, err
}
return casUpdate(kv, key, mutate, false)
}
return nil, err
}
return updated, nil
}
// ListTerminatedInstances returns all terminated instances from the terminated KV bucket.
func (m *JetStreamManager) ListTerminatedInstances() ([]*vm.VM, error) {
if m.terminatedKV == nil {
+198
View File
@@ -456,6 +456,103 @@ func TestJetStreamManager_DeleteStoppedInstance(t *testing.T) {
require.NoError(t, err)
}
// TestJetStreamManager_ClaimStoppedInstance_HappyPath verifies a claim
// atomically removes the record and hands back the VM it held.
func TestJetStreamManager_ClaimStoppedInstance_HappyPath(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
testVM := &vm.VM{ID: "i-claim-happy", Status: vm.StateStopped, InstanceType: "t3.micro"}
require.NoError(t, jsm.WriteStoppedInstance(testVM.ID, testVM))
claimed, err := jsm.ClaimStoppedInstance(testVM.ID)
require.NoError(t, err)
require.NotNil(t, claimed)
assert.Equal(t, testVM.ID, claimed.ID)
assert.Equal(t, "t3.micro", claimed.InstanceType)
// The claim must have removed the record — a second claim, or a plain
// Load, must both observe it gone.
loaded, err := jsm.LoadStoppedInstance(testVM.ID)
require.NoError(t, err)
assert.Nil(t, loaded, "a successful claim must remove the record")
}
// TestJetStreamManager_ClaimStoppedInstance_NotFound verifies claiming an id
// with no record returns vm.ErrStoppedInstanceClaimed rather than a bare
// nats not-found error, so callers can map it directly to the API error.
func TestJetStreamManager_ClaimStoppedInstance_NotFound(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
_, err = jsm.ClaimStoppedInstance("i-does-not-exist")
assert.ErrorIs(t, err, vm.ErrStoppedInstanceClaimed)
}
// TestJetStreamManager_ClaimStoppedInstance_ConcurrentClaim is the
// lower-level regression test for the double-start bug: it fires N
// concurrent ClaimStoppedInstance calls at the same key against a real
// JetStream KV bucket and asserts exactly one succeeds — proving the
// nats.LastRevision-guarded Delete this claim is built on is a genuine
// mutual-exclusion primitive under real NATS wire semantics, not just at
// the fake/mock layer exercised by the handlers-level test.
func TestJetStreamManager_ClaimStoppedInstance_ConcurrentClaim(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
testVM := &vm.VM{ID: "i-claim-race", Status: vm.StateStopped, InstanceType: "t3.micro"}
require.NoError(t, jsm.WriteStoppedInstance(testVM.ID, testVM))
const n = 8
results := make([]*vm.VM, n)
errs := make([]error, n)
var wg sync.WaitGroup
wg.Add(n)
for i := range n {
go func(i int) {
defer wg.Done()
results[i], errs[i] = jsm.ClaimStoppedInstance(testVM.ID)
}(i)
}
wg.Wait()
var won, lost int
for i := range n {
switch {
case errs[i] == nil:
won++
require.NotNil(t, results[i])
assert.Equal(t, testVM.ID, results[i].ID)
case errors.Is(errs[i], vm.ErrStoppedInstanceClaimed):
lost++
default:
t.Fatalf("unexpected error from ClaimStoppedInstance: %v", errs[i])
}
}
assert.Equal(t, 1, won, "exactly one concurrent claim must succeed")
assert.Equal(t, n-1, lost, "every other concurrent claim must lose with ErrStoppedInstanceClaimed")
loaded, err := jsm.LoadStoppedInstance(testVM.ID)
require.NoError(t, err)
assert.Nil(t, loaded, "the winning claim must have removed the record")
}
// TestJetStreamManager_ListStoppedInstances tests listing multiple stopped instances
func TestJetStreamManager_ListStoppedInstances(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
@@ -1113,6 +1210,107 @@ func TestJetStreamManager_TerminatedInstance_KVNotInitialized(t *testing.T) {
assert.Error(t, err)
}
// --- CAS conflict-retry tests ---
// TestJetStreamManager_UpdateTerminatedInstance_ConflictRetry locks the CAS
// contract for the multi-writer Teardown map: a concurrent writer that lands
// between UpdateTerminatedInstance's Get and Update must be detected via a
// revision conflict and retried, and BOTH updates — the injected concurrent
// one and the retried one — must land, none silently dropped.
func TestJetStreamManager_UpdateTerminatedInstance_ConflictRetry(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
testVM := &vm.VM{ID: "i-cas-conflict", Teardown: map[string]string{"volumes": "pending", "gpu": "pending"}}
require.NoError(t, jsm.WriteTerminatedInstance(testVM.ID, testVM))
defer func() { _ = jsm.DeleteTerminatedInstance(testVM.ID) }()
var attempts int
var injectOnce sync.Once
updated, err := jsm.UpdateTerminatedInstance(testVM.ID, func(v *vm.VM) {
attempts++
if v.Teardown == nil {
v.Teardown = map[string]string{}
}
v.Teardown["gpu"] = "done"
// Simulate a concurrent writer landing between our Get and Update on
// the first attempt only, forcing exactly one revision conflict.
injectOnce.Do(func() {
_, cerr := jsm.UpdateTerminatedInstance(testVM.ID, func(cv *vm.VM) {
if cv.Teardown == nil {
cv.Teardown = map[string]string{}
}
cv.Teardown["volumes"] = "done"
})
require.NoError(t, cerr)
})
})
require.NoError(t, err, "the retried Update must eventually succeed against the fresh revision")
require.NotNil(t, updated)
assert.Equal(t, 2, attempts, "the injected concurrent write must force exactly one retry")
assert.Equal(t, "done", updated.Teardown["gpu"])
loaded, err := jsm.LoadTerminatedInstance(testVM.ID)
require.NoError(t, err)
require.NotNil(t, loaded)
assert.Equal(t, "done", loaded.Teardown["volumes"], "the concurrent writer's update must survive, not be clobbered")
assert.Equal(t, "done", loaded.Teardown["gpu"], "the retried update must also land")
}
// TestJetStreamManager_UpdateTerminatedInstance_NotFound verifies
// UpdateTerminatedInstance surfaces a clear error rather than silently
// creating a record when there is nothing to merge into.
func TestJetStreamManager_UpdateTerminatedInstance_NotFound(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
require.NoError(t, jsm.InitTerminatedInstanceBucket())
_, err = jsm.UpdateTerminatedInstance("i-does-not-exist", func(v *vm.VM) {})
assert.ErrorIs(t, err, nats.ErrKeyNotFound)
}
// TestJetStreamManager_WriteStoppedInstance_OverwritesConcurrentValue verifies
// WriteStoppedInstance's CAS write still succeeds and lands the caller's full
// snapshot even when a concurrent writer changed the key first — the
// conflict-retry loop must not surface a spurious error for the common,
// non-merging "replace wholesale" write path.
func TestJetStreamManager_WriteStoppedInstance_OverwritesConcurrentValue(t *testing.T) {
nc, err := nats.Connect(sharedJSNATSURL)
require.NoError(t, err)
defer nc.Close()
jsm, err := NewJetStreamManager(nc, 1)
require.NoError(t, err)
require.NoError(t, jsm.InitKVBucket())
first := &vm.VM{ID: "i-stopped-cas", InstanceType: "t3.micro"}
require.NoError(t, jsm.WriteStoppedInstance(first.ID, first))
defer func() { _ = jsm.DeleteStoppedInstance(first.ID) }()
// A concurrent writer changes the record between the caller's decision to
// write and WriteStoppedInstance's internal Get.
concurrent := &vm.VM{ID: "i-stopped-cas", InstanceType: "t3.small"}
require.NoError(t, jsm.WriteStoppedInstance(concurrent.ID, concurrent))
final := &vm.VM{ID: "i-stopped-cas", InstanceType: "t3.large"}
require.NoError(t, jsm.WriteStoppedInstance(final.ID, final))
loaded, err := jsm.LoadStoppedInstance(final.ID)
require.NoError(t, err)
require.NotNil(t, loaded)
assert.Equal(t, "t3.large", loaded.InstanceType)
}
// --- Terminated KV bucket recovery tests ---
// deleteTerminatedInstanceBucket deletes the underlying JetStream stream for the terminated-instances KV bucket.
+9
View File
@@ -52,6 +52,10 @@ func (a *stateStoreAdapter) DeleteStoppedInstance(id string) error {
return a.js.DeleteStoppedInstance(id)
}
func (a *stateStoreAdapter) ClaimStoppedInstance(id string) (*vm.VM, error) {
return a.js.ClaimStoppedInstance(id)
}
func (a *stateStoreAdapter) ListStoppedInstances() ([]*vm.VM, error) {
return a.js.ListStoppedInstances()
}
@@ -60,6 +64,10 @@ func (a *stateStoreAdapter) WriteTerminatedInstance(id string, v *vm.VM) error {
return a.js.WriteTerminatedInstance(id, v)
}
func (a *stateStoreAdapter) UpdateTerminatedInstance(id string, mutate func(*vm.VM)) (*vm.VM, error) {
return a.js.UpdateTerminatedInstance(id, mutate)
}
func (a *stateStoreAdapter) ListTerminatedInstances() ([]*vm.VM, error) {
return a.js.ListTerminatedInstances()
}
@@ -550,6 +558,7 @@ func (d *Daemon) buildVMManagerDeps() vm.Deps {
DevNetworking: d.config.Daemon.DevNetworking,
BindHost: d.config.Host,
DetachDelay: d.detachDelay,
DeviceDeletedTimeout: d.deviceDeletedTimeout,
ConsumeCleanShutdownMarker: d.consumeCleanShutdownMarker(),
}
}
+1 -1
View File
@@ -22,7 +22,7 @@ func seedDiscoveryCluster(t *testing.T) http.Handler {
t.Helper()
_, nc, js := testutil.StartTestJetStream(t)
kv, err := handlers_eks.GetOrCreateAccountBucket(js, testDiscAccount)
kv, err := handlers_eks.GetOrCreateAccountBucket(js, testDiscAccount, 1)
if err != nil {
t.Fatalf("account bucket: %v", err)
}
@@ -86,6 +86,11 @@ type StoppedInstanceStore interface {
WriteStoppedInstance(instanceID string, instance *vm.VM) error
DeleteStoppedInstance(instanceID string) error
WriteTerminatedInstance(instanceID string, instance *vm.VM) error
// ClaimStoppedInstance atomically removes instanceID's record and
// returns the VM it held, so at most one caller can ever win a race to
// (re)launch the same stopped instance. Returns vm.ErrStoppedInstanceClaimed
// if a concurrent caller already claimed (or otherwise removed) the record.
ClaimStoppedInstance(instanceID string) (*vm.VM, error)
}
// InstanceTagWriter projects an instance record's full tag set into the
+37 -11
View File
@@ -2161,6 +2161,25 @@ func (s *InstanceServiceImpl) StartStoppedInstance(ctx context.Context, input *S
return nil, errors.New(awserrors.ErrorInvalidInstanceIDNotFound)
}
// Atomic claim: at most one caller across the cluster can win this race
// — a forwarded call and this node's local fallback, or two nodes
// racing the same forwarded call, both pass the checks above, but only
// one succeeds here. The loser must not touch resourceMgr/vmMgr at
// all, closing the double-start-onto-the-same-volume race. This also
// replaces the plain Load used above as the record of truth going
// forward: `instance` is reassigned to the freshest claimed value.
instance, err = s.stoppedStore.ClaimStoppedInstance(input.InstanceID)
if err != nil {
if errors.Is(err, vm.ErrStoppedInstanceClaimed) {
slog.WarnContext(ctx, "StartStoppedInstance: lost race to claim stopped instance",
"instanceId", input.InstanceID)
return nil, errors.New(awserrors.ErrorIncorrectInstanceState)
}
slog.ErrorContext(ctx, "StartStoppedInstance: failed to claim stopped instance",
"instanceId", input.InstanceID, "err", err)
return nil, errors.New(awserrors.ErrorServerInternal)
}
// Reset node-local fields that are stale after cross-node migration.
instance.ResetNodeLocalState()
@@ -2168,10 +2187,12 @@ func (s *InstanceServiceImpl) StartStoppedInstance(ctx context.Context, input *S
if !ok {
slog.ErrorContext(ctx, "StartStoppedInstance: instance type not available on this node",
"instanceId", input.InstanceID, "instanceType", instance.InstanceType)
s.restoreClaimedStoppedInstance(ctx, instance)
return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity)
}
if err := s.resourceMgr.Allocate(instanceType); err != nil {
slog.ErrorContext(ctx, "StartStoppedInstance: failed to allocate resources", "instanceId", input.InstanceID, "err", err)
s.restoreClaimedStoppedInstance(ctx, instance)
return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity)
}
@@ -2188,6 +2209,7 @@ func (s *InstanceServiceImpl) StartStoppedInstance(ctx context.Context, input *S
slog.ErrorContext(ctx, "StartStoppedInstance: GPU claim failed", "instanceId", input.InstanceID, "err", gpuErr)
s.resourceMgr.Deallocate(instanceType)
s.vmMgr.Delete(instance.ID)
s.restoreClaimedStoppedInstance(ctx, instance)
return nil, errors.New(awserrors.ErrorInsufficientInstanceCapacity)
}
instance.GPUAttachments = []gpu.GPUAttachment{*att}
@@ -2206,27 +2228,31 @@ func (s *InstanceServiceImpl) StartStoppedInstance(ctx context.Context, input *S
}
s.resourceMgr.Deallocate(instanceType)
s.vmMgr.Delete(instance.ID)
s.restoreClaimedStoppedInstance(ctx, instance)
return nil, errors.New(awserrors.ErrorServerInternal)
}
// Discover actual guest device names via QMP query-block.
s.vmMgr.UpdateGuestDeviceNames(instance)
// Remove from shared KV now that it's running locally. Retry once on failure —
// a stale KV entry risks duplicate starts.
if err := s.stoppedStore.DeleteStoppedInstance(input.InstanceID); err != nil {
slog.WarnContext(ctx, "StartStoppedInstance: first KV delete failed, retrying",
"instanceId", input.InstanceID, "err", err)
if retryErr := s.stoppedStore.DeleteStoppedInstance(input.InstanceID); retryErr != nil {
slog.ErrorContext(ctx, "StartStoppedInstance: KV delete failed after retry, instance is running locally but stale entry remains in shared KV",
"instanceId", input.InstanceID, "err", retryErr)
}
}
slog.InfoContext(ctx, "Started stopped instance from shared KV", "instanceId", instance.ID)
return &StartStoppedInstanceOutput{Status: "running", InstanceID: instance.ID}, nil
}
// restoreClaimedStoppedInstance writes instance back to the shared
// stopped-instance store after a downstream failure that follows a
// successful ClaimStoppedInstance (which already removed it), so the claim
// does not silently lose the instance. Best-effort: the caller already has
// an error to return, so a restore failure is logged, not surfaced — the
// instance is still safely stopped in the local vm.Manager's view (it was
// never inserted, or was removed again above).
func (s *InstanceServiceImpl) restoreClaimedStoppedInstance(ctx context.Context, instance *vm.VM) {
if err := s.stoppedStore.WriteStoppedInstance(instance.ID, instance); err != nil {
slog.ErrorContext(ctx, "StartStoppedInstance: failed to restore claimed instance to shared KV after downstream failure",
"instanceId", instance.ID, "err", err)
}
}
// TerminateStoppedInstance terminates a stopped instance: deletes its volumes,
// releases its public IP and ENI, and moves it to the terminated KV bucket.
func (s *InstanceServiceImpl) TerminateStoppedInstance(ctx context.Context, input *TerminateStoppedInstanceInput, accountID string) (*TerminateStoppedInstanceOutput, error) {
@@ -6,6 +6,7 @@ import (
"errors"
"fmt"
"strings"
"sync"
"testing"
"time"
@@ -516,6 +517,9 @@ func TestParseVolumeParams_Io1WithIops(t *testing.T) {
// --- Describe* coverage -----------------------------------------------------
type fakeResourceCapacityProvider struct {
// mu guards the mutating fields below so concurrent-claim tests can
// call Allocate/Deallocate from multiple goroutines race-free.
mu sync.Mutex
types []*ec2.InstanceTypeInfo
supportedTypes []*ec2.InstanceTypeInfo
gotShowCap bool
@@ -552,6 +556,8 @@ func (f *fakeResourceCapacityProvider) GetSupportedInstanceTypeInfos() []*ec2.In
}
func (f *fakeResourceCapacityProvider) Allocate(it *ec2.InstanceTypeInfo) error {
f.mu.Lock()
defer f.mu.Unlock()
if f.allocateErr != nil {
return f.allocateErr
}
@@ -560,6 +566,8 @@ func (f *fakeResourceCapacityProvider) Allocate(it *ec2.InstanceTypeInfo) error
}
func (f *fakeResourceCapacityProvider) Deallocate(it *ec2.InstanceTypeInfo) {
f.mu.Lock()
defer f.mu.Unlock()
f.deallocated = append(f.deallocated, it)
}
@@ -594,6 +602,7 @@ func (f *fakeResourceCapacityProvider) InstanceTypes() map[string]*ec2.InstanceT
}
type fakeStoppedStore struct {
mu sync.Mutex
stopped []*vm.VM
terminated []*vm.VM
loadByID map[string]*vm.VM
@@ -608,12 +617,22 @@ type fakeStoppedStore struct {
deleteErr error
deleteFailFirst bool
deleteAttempts int
// claimedStopped / claimErr drive ClaimStoppedInstance, the atomic
// claim used by StartStoppedInstance. Claiming removes the entry from
// loadByID (mirroring the real KV delete-as-claim), guarded by mu so
// concurrent StartStoppedInstance callers race exactly like production.
claimedStopped []string
claimErr error
claimAttempts int
}
func (f *fakeStoppedStore) LoadStoppedInstance(id string) (*vm.VM, error) {
if f.loadErr != nil {
return nil, f.loadErr
}
f.mu.Lock()
defer f.mu.Unlock()
if v, ok := f.loadByID[id]; ok {
return v, nil
}
@@ -623,18 +642,24 @@ func (f *fakeStoppedStore) ListStoppedInstances() ([]*vm.VM, error) {
if f.listErr != nil {
return nil, f.listErr
}
f.mu.Lock()
defer f.mu.Unlock()
return f.stopped, nil
}
func (f *fakeStoppedStore) ListTerminatedInstances() ([]*vm.VM, error) {
if f.listTermErr != nil {
return nil, f.listTermErr
}
f.mu.Lock()
defer f.mu.Unlock()
return f.terminated, nil
}
func (f *fakeStoppedStore) WriteStoppedInstance(id string, instance *vm.VM) error {
if f.writeErr != nil {
return f.writeErr
}
f.mu.Lock()
defer f.mu.Unlock()
if f.wroteStopped == nil {
f.wroteStopped = make(map[string]*vm.VM)
}
@@ -645,6 +670,8 @@ func (f *fakeStoppedStore) WriteTerminatedInstance(id string, instance *vm.VM) e
if f.writeTermErr != nil {
return f.writeTermErr
}
f.mu.Lock()
defer f.mu.Unlock()
if f.wroteTerminated == nil {
f.wroteTerminated = make(map[string]*vm.VM)
}
@@ -652,18 +679,44 @@ func (f *fakeStoppedStore) WriteTerminatedInstance(id string, instance *vm.VM) e
return nil
}
func (f *fakeStoppedStore) DeleteStoppedInstance(id string) error {
f.mu.Lock()
f.deleteAttempts++
if f.deleteFailFirst && f.deleteAttempts == 1 {
attempt := f.deleteAttempts
f.mu.Unlock()
if f.deleteFailFirst && attempt == 1 {
return errors.New("transient delete failure")
}
if f.deleteErr != nil {
return f.deleteErr
}
f.mu.Lock()
defer f.mu.Unlock()
f.deletedStopped = append(f.deletedStopped, id)
return nil
}
func (f *fakeStoppedStore) DeleteTerminatedInstance(string) error { return nil }
// ClaimStoppedInstance mimics the real atomic delete-as-claim: under the
// store lock, remove and return loadByID[id], or vm.ErrStoppedInstanceClaimed
// if it is already gone — mirroring the real KV's "revision conflict or
// not-found" outcome for a losing racer. The lock makes this genuinely
// exclusive under concurrent callers, matching the production guarantee.
func (f *fakeStoppedStore) ClaimStoppedInstance(id string) (*vm.VM, error) {
f.mu.Lock()
defer f.mu.Unlock()
f.claimAttempts++
if f.claimErr != nil {
return nil, f.claimErr
}
v, ok := f.loadByID[id]
if !ok {
return nil, vm.ErrStoppedInstanceClaimed
}
delete(f.loadByID, id)
f.claimedStopped = append(f.claimedStopped, id)
return v, nil
}
func TestDescribeInstanceTypes_NilResourceMgr(t *testing.T) {
svc := &InstanceServiceImpl{}
_, err := svc.DescribeInstanceTypes(context.Background(), &ec2.DescribeInstanceTypesInput{}, "")
@@ -1942,6 +1995,12 @@ func TestStartStoppedInstance_InstanceTypeUnknown(t *testing.T) {
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error())
assert.Empty(t, prov.allocated, "no allocation should occur for unknown type")
// The claim removed the entry from KV before the instance-type check;
// the rollback must write it back so the instance is not lost.
assert.Contains(t, store.claimedStopped, id, "claim must run before the instance-type check")
require.NotNil(t, store.wroteStopped[id], "claimed instance must be restored to KV after a downstream failure")
assert.Equal(t, vm.StateStopped, store.wroteStopped[id].Status)
}
func TestStartStoppedInstance_AllocateFails(t *testing.T) {
@@ -1964,6 +2023,12 @@ func TestStartStoppedInstance_AllocateFails(t *testing.T) {
_, err := svc.StartStoppedInstance(context.Background(), &StartStoppedInstanceInput{InstanceID: id}, "acc")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorInsufficientInstanceCapacity, err.Error())
// A failed Allocate must not leave the instance stranded — it was
// already claimed (removed from KV) by this point, so it must be
// restored.
require.NotNil(t, store.wroteStopped[id], "claimed instance must be restored to KV after Allocate failure")
assert.Equal(t, vm.StateStopped, store.wroteStopped[id].Status)
}
// GPU claim failure must roll back the resource allocation and remove the VM
@@ -1998,6 +2063,132 @@ func TestStartStoppedInstance_GPUClaimFailureRollsBack(t *testing.T) {
_, stillInMgr := mgr.Get(id)
assert.False(t, stillInMgr, "GPU claim failure must remove the VM from the manager map")
assert.Empty(t, store.deletedStopped, "stopped-KV entry must remain on rollback")
// The atomic claim removed the entry from KV before Allocate/GPU-claim
// ran; the rollback must write it back so it is not lost.
require.NotNil(t, store.wroteStopped[id], "claimed instance must be restored to KV after GPU claim failure")
assert.Equal(t, vm.StateStopped, store.wroteStopped[id].Status)
}
// TestStartStoppedInstance_ClaimConflict proves a lost claim race is
// reported as ErrorIncorrectInstanceState and never reaches Allocate/Insert.
func TestStartStoppedInstance_ClaimConflict(t *testing.T) {
id := "i-claim-conflict"
itype := "t3.micro"
store := &fakeStoppedStore{
loadByID: map[string]*vm.VM{
id: {ID: id, Status: vm.StateStopped, AccountID: "acc", InstanceType: itype},
},
claimErr: vm.ErrStoppedInstanceClaimed,
}
prov := &fakeResourceCapacityProvider{
instanceTypes: map[string]*ec2.InstanceTypeInfo{itype: {InstanceType: aws.String(itype)}},
}
svc := &InstanceServiceImpl{
stoppedStore: store,
resourceMgr: prov,
vmMgr: vm.NewManager(),
}
_, err := svc.StartStoppedInstance(context.Background(), &StartStoppedInstanceInput{InstanceID: id}, "acc")
require.Error(t, err)
assert.Equal(t, awserrors.ErrorIncorrectInstanceState, err.Error())
assert.Empty(t, prov.allocated, "a lost claim race must never reach Allocate")
}
// raceVolumeMounter is a no-op vm.VolumeMounter that lets vmMgr.Run proceed
// past the mount step in the concurrent-claim race test below, so Run
// reaches vm.Manager's (deliberately unwired) InstanceTypeResolver check and
// fails fast and deterministically without touching a real qemu process.
type raceVolumeMounter struct{}
func (raceVolumeMounter) Mount(*vm.VM) error { return nil }
func (raceVolumeMounter) Unmount(*vm.VM) error { return nil }
func (raceVolumeMounter) MountOne(*spxtypes.EBSRequest) error { return nil }
func (raceVolumeMounter) UnmountOne(spxtypes.EBSRequest) error { return nil }
// TestStartStoppedInstance_ConcurrentClaimRace is the regression test for
// the double-start bug this claim closes: two nodes (or a forwarded call
// racing a local fallback) could both observe the same stopped instance as
// claimable and both launch it onto the same viperblock volume. It fires two
// concurrent StartStoppedInstance calls at the same stopped instance id and
// asserts the atomic claim lets exactly one of them proceed past it —
// reaching resourceMgr.Allocate / vmMgr.Insert / vmMgr.Run — while the other
// loses the race and never touches resourceMgr or vmMgr at all.
//
// The loser's error code depends on exactly when it lost: if it loses the
// atomic ClaimStoppedInstance call itself, it gets ErrorIncorrectInstanceState;
// if the winner's claim already completed before the loser's preliminary
// (non-atomic, pre-claim) LoadStoppedInstance validation ran, the loser sees
// the record as already gone and gets ErrorInvalidInstanceIDNotFound instead.
// Both are safe, non-destructive outcomes — the property under test is that
// exactly one caller ever reaches past the claim, not which of the two
// equally-valid error codes the loser receives.
//
// The winner's vmMgr.Run is wired to fail fast (no InstanceTypeResolver) so
// the test stays deterministic without exec'ing a real qemu process; that
// failure also exercises restoreClaimedStoppedInstance, so this test doubles
// as coverage that a downstream failure after a successful claim does not
// lose the instance.
func TestStartStoppedInstance_ConcurrentClaimRace(t *testing.T) {
t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
id := "i-race"
itype := "t3.micro"
store := &fakeStoppedStore{loadByID: map[string]*vm.VM{
id: {ID: id, Status: vm.StateStopped, AccountID: "acc", InstanceType: itype},
}}
prov := &fakeResourceCapacityProvider{
instanceTypes: map[string]*ec2.InstanceTypeInfo{itype: {InstanceType: aws.String(itype)}},
}
mgr := vm.NewManagerWithDeps(vm.Deps{VolumeMounter: raceVolumeMounter{}})
svc := &InstanceServiceImpl{
stoppedStore: store,
resourceMgr: prov,
vmMgr: mgr,
}
const n = 2
errs := make([]error, n)
var wg sync.WaitGroup
wg.Add(n)
for i := range n {
go func(i int) {
defer wg.Done()
_, err := svc.StartStoppedInstance(context.Background(), &StartStoppedInstanceInput{InstanceID: id}, "acc")
errs[i] = err
}(i)
}
wg.Wait()
var conflicted, winnerFailed int
for _, err := range errs {
require.Error(t, err, "vmMgr.Run is wired to fail fast — no caller should succeed end-to-end")
switch err.Error() {
case awserrors.ErrorIncorrectInstanceState, awserrors.ErrorInvalidInstanceIDNotFound:
conflicted++
case awserrors.ErrorServerInternal:
winnerFailed++
default:
t.Fatalf("unexpected error: %v", err)
}
}
assert.Equal(t, 1, conflicted, "exactly one caller must lose the claim race")
assert.Equal(t, 1, winnerFailed, "exactly one caller must win the claim and proceed to vmMgr.Run")
assert.Len(t, prov.allocated, 1, "only the claim winner may reach resourceMgr.Allocate")
// Exactly one caller ever wins the atomic claim, regardless of whether
// the loser lost at the claim itself or earlier at the preliminary Load
// (see comment above) — it may not always attempt the claim at all.
assert.Len(t, store.claimedStopped, 1, "the atomic claim must succeed exactly once")
assert.GreaterOrEqual(t, store.claimAttempts, 1, "the winner must have attempted the claim")
require.NotNil(t, store.wroteStopped[id], "the winner's downstream Run failure must restore the instance to KV")
assert.Equal(t, vm.StateStopped, store.wroteStopped[id].Status)
_, stillInMgr := mgr.Get(id)
assert.False(t, stillInMgr, "a failed Run must not leave the VM in the manager map")
}
// --- PrepareRunInstances / ec2.cmd dispatch tests ---------------------------
+7 -5
View File
@@ -35,14 +35,16 @@ type StagedAddonManifest struct {
// The VM-side delivery slice applies it via the K3s auto-deploy dir; until then
// the add-on sits CREATING.
type stagingInstaller struct {
nc *nats.Conn
nc *nats.Conn
clusterSize int
}
var _ AddonInstaller = (*stagingInstaller)(nil)
// newStagingInstaller returns a stagingInstaller bound to the daemon NATS connection.
func newStagingInstaller(nc *nats.Conn) *stagingInstaller {
return &stagingInstaller{nc: nc}
// newStagingInstaller returns a stagingInstaller bound to the daemon NATS
// connection, using clusterSize as the per-account bucket's replica count.
func newStagingInstaller(nc *nats.Conn, clusterSize int) *stagingInstaller {
return &stagingInstaller{nc: nc, clusterSize: clusterSize}
}
func (i *stagingInstaller) acctKV(accountID string) (nats.KeyValue, error) {
@@ -53,7 +55,7 @@ func (i *stagingInstaller) acctKV(accountID string) (nats.KeyValue, error) {
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
return GetOrCreateAccountBucket(js, accountID)
return GetOrCreateAccountBucket(js, accountID, max(i.clusterSize, 1))
}
func (i *stagingInstaller) Install(accountID, cluster string, rec *AddonRecord) error {
@@ -14,7 +14,7 @@ func acctKVForTest(t *testing.T, svc *EKSServiceImpl) nats.KeyValue {
t.Helper()
js, err := svc.deps.NATSConn.JetStream()
require.NoError(t, err)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
return kv
}
+1 -1
View File
@@ -329,7 +329,7 @@ func (s *EKSServiceImpl) addonInstaller() AddonInstaller {
if s.deps.AddonInstaller != nil {
return s.deps.AddonInstaller
}
return newStagingInstaller(s.deps.NATSConn)
return newStagingInstaller(s.deps.NATSConn, s.deps.ClusterSize)
}
// markAddonFailed best-effort flips a record to CREATE_FAILED with the error reason.
+1 -1
View File
@@ -243,7 +243,7 @@ func TestStagingInstaller_StagesManifest(t *testing.T) {
js, err := svc.deps.NATSConn.JetStream()
require.NoError(t, err)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
entry, err := kv.Get(AddonManifestKey("c1", albController))
require.NoError(t, err, "installer must stage a manifest for VM-side delivery")
@@ -19,9 +19,9 @@ import (
func newStateReconcilerHarness(t *testing.T, opts ...ReconcilerOption) (*ClusterReconciler, *nats.Conn, nats.KeyValue) {
t.Helper()
_, nc, js := testutil.StartTestJetStream(t)
leaderKV, err := InitLeaderBucket(js)
leaderKV, err := InitLeaderBucket(js, 1)
require.NoError(t, err)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(acctKV, sampleClusterMeta("alpha")))
@@ -48,9 +48,9 @@ func (s *stubHTTPDoer) setResponse(status int, err error) {
func newReconcilerHarness(t *testing.T, healthURL string, opts ...ReconcilerOption) (*ClusterReconciler, nats.KeyValue, nats.KeyValue) {
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
leaderKV, err := InitLeaderBucket(js)
leaderKV, err := InitLeaderBucket(js, 1)
require.NoError(t, err)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(acctKV, sampleClusterMeta("alpha")))
@@ -61,9 +61,9 @@ func newReconcilerHarness(t *testing.T, healthURL string, opts ...ReconcilerOpti
func TestNewClusterReconciler_EmptyInputsRejected(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
leaderKV, err := InitLeaderBucket(js)
leaderKV, err := InitLeaderBucket(js, 1)
require.NoError(t, err)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
_, err = NewClusterReconciler(nil, acctKV, testAccountID, "alpha", "h", "")
+1 -1
View File
@@ -15,7 +15,7 @@ import (
func newClusterStateTestKV(t *testing.T) nats.KeyValue {
t.Helper()
_, _, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
return kv
}
+1 -1
View File
@@ -57,7 +57,7 @@ type eksServiceFixture struct {
func newEKSServiceFixture(t *testing.T) *eksServiceFixture {
t.Helper()
_, nc, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
nlb := newFakeNLBProvisioner()
+1 -1
View File
@@ -68,7 +68,7 @@ func (r *EKSBillableReaper) Sweep(ctx context.Context) (int, error) {
continue // ENI gone or untagged: cannot confirm orphan-hood, never reap
}
acctKV, err := GetOrCreateAccountBucket(js, clusterAccount)
acctKV, err := GetOrCreateAccountBucket(js, clusterAccount, max(r.svc.deps.ClusterSize, 1))
if err != nil {
slog.WarnContext(ctx, "eks-billable: account bucket lookup failed", "account", clusterAccount, "err", err)
continue
+2 -2
View File
@@ -37,7 +37,7 @@ type natsBootstrapHarness struct {
func newBootstrapHarness(t *testing.T) *natsBootstrapHarness {
t.Helper()
_, nc, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(kv, sampleClusterMeta(bootstrapTestCluster)))
@@ -82,7 +82,7 @@ func (h *natsBootstrapHarness) waitSubsBound(base, want int) {
func TestNewNATSBootstrap_RejectsBadInputs(t *testing.T) {
_, nc, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
_, err = NewNATSBootstrap(nil, kv, bootstrapTestMasterKey, testAccountID, "alpha")
+1 -1
View File
@@ -163,7 +163,7 @@ func (s *EKSServiceImpl) nodegroupAcctKV(accountID string) (nats.KeyValue, error
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
+1 -1
View File
@@ -175,7 +175,7 @@ func TestReclaimOrphanedNodegroups_TerminatesStrandedWorkers(t *testing.T) {
func TestNodegroupRecord_CRUDRoundTrip(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
rec := &NodegroupRecord{
@@ -12,7 +12,7 @@ import (
func TestLoadRecoveryDirective_AbsentIsZeroNone(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
d, err := LoadRecoveryDirective(acctKV, "alpha", "i-missing")
@@ -23,7 +23,7 @@ func TestLoadRecoveryDirective_AbsentIsZeroNone(t *testing.T) {
func TestStoreRecoveryDirective_EpochIncrements(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID)
acctKV, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
first, err := StoreRecoveryDirective(acctKV, "alpha", "i-0", RecoveryActionClusterReset, "")
+13 -8
View File
@@ -44,6 +44,11 @@ type EKSServiceDeps struct {
Region string
HolderID string
// ClusterSize is the daemon's node count, used as the JetStream replica
// count for the lazily-created per-account and leader KV buckets so they
// match the cluster's other R3 streams instead of staying stuck at R1.
ClusterSize int
// InternalSuffix is the AWS-parity internal DNS suffix (e.g. spinifex.internal)
// used to compose the worker's ECR registry host.
InternalSuffix string
@@ -244,7 +249,7 @@ func NewEKSServiceImpl(deps EKSServiceDeps) (*EKSServiceImpl, error) {
if err != nil {
return nil, fmt.Errorf("failed to get JetStream context: %w", err)
}
leaderKV, err := InitLeaderBucket(js)
leaderKV, err := InitLeaderBucket(js, max(deps.ClusterSize, 1))
if err != nil {
return nil, err
}
@@ -417,7 +422,7 @@ func (s *EKSServiceImpl) CreateCluster(ctx context.Context, input *eks.CreateClu
if err != nil {
return nil, logCreateErr(name, accountID, "jetstream", err)
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
return nil, logCreateErr(name, accountID, "get account bucket", err)
}
@@ -843,7 +848,7 @@ func (s *EKSServiceImpl) DescribeCluster(ctx context.Context, input *eks.Describ
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
@@ -862,7 +867,7 @@ func (s *EKSServiceImpl) ListClusters(ctx context.Context, input *eks.ListCluste
if err != nil {
return nil, eksReadUnavailableOr(err, "jetstream")
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
return nil, eksReadUnavailableOr(err, "get account bucket")
}
@@ -909,7 +914,7 @@ func (s *EKSServiceImpl) DeleteCluster(ctx context.Context, input *eks.DeleteClu
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
@@ -1296,7 +1301,7 @@ func (s *EKSServiceImpl) acctKVForCluster(accountID, cluster string) (nats.KeyVa
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
return nil, fmt.Errorf("get account bucket: %w", err)
}
@@ -1758,7 +1763,7 @@ func (s *EKSServiceImpl) accountBucket(accountID string) (nats.KeyValue, error)
if err != nil {
return nil, fmt.Errorf("jetstream: %w", err)
}
return GetOrCreateAccountBucket(js, accountID)
return GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
}
// clusterNameFromARN extracts the cluster name from an EKS cluster ARN
@@ -1931,7 +1936,7 @@ func (s *EKSServiceImpl) spawnReconciler(accountID, clusterName string, _ *Clust
slog.Error("spawnReconciler: jetstream", "err", err)
return
}
acctKV, err := GetOrCreateAccountBucket(js, accountID)
acctKV, err := GetOrCreateAccountBucket(js, accountID, max(s.deps.ClusterSize, 1))
if err != nil {
slog.Error("spawnReconciler: account bucket", "err", err)
return
+3 -3
View File
@@ -195,7 +195,7 @@ func seedTestCluster(t *testing.T, svc *EKSServiceImpl, cluster string) {
t.Helper()
js, err := svc.deps.NATSConn.JetStream()
require.NoError(t, err)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
require.NoError(t, PutClusterMeta(kv, &ClusterMeta{Name: cluster, Status: ClusterStatusActive}))
}
@@ -374,7 +374,7 @@ func TestEKSServiceImpl_ClusterTagRoundTrip(t *testing.T) {
svc := setupTestService(t)
js, err := svc.deps.NATSConn.JetStream()
require.NoError(t, err)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
const arn = "arn:aws:eks:us-east-1:111122223333:cluster/c1"
@@ -417,7 +417,7 @@ func TestEKSServiceImpl_NodegroupTagRoundTrip(t *testing.T) {
svc := setupTestService(t)
js, err := svc.deps.NATSConn.JetStream()
require.NoError(t, err)
kv, err := GetOrCreateAccountBucket(js, testAccountID)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
const arn = "arn:aws:eks:us-east-1:111122223333:nodegroup/c1/ng1/abc123"
+16 -13
View File
@@ -181,11 +181,12 @@ func AccountBucketName(accountID string) string {
}
// GetOrCreateAccountBucket returns the per-account KV bucket for accountID,
// creating it on first use. Idempotent: subsequent calls with the same
// accountID return the existing handle.
func GetOrCreateAccountBucket(js nats.JetStreamContext, accountID string) (nats.KeyValue, error) {
// creating it on first use at the given replica count (clamped to a minimum
// of 1). Idempotent: subsequent calls with the same accountID return the
// existing handle.
func GetOrCreateAccountBucket(js nats.JetStreamContext, accountID string, replicas int) (nats.KeyValue, error) {
bucket := AccountBucketName(accountID)
kv, err := utils.GetOrCreateKVBucket(js, bucket, KVBucketEKSAccountHistory)
kv, err := utils.GetOrCreateKVBucketWithReplicas(js, bucket, KVBucketEKSAccountHistory, replicas)
if err != nil {
return nil, fmt.Errorf("failed to create EKS per-account KV bucket %s: %w", bucket, err)
}
@@ -196,16 +197,18 @@ func GetOrCreateAccountBucket(js nats.JetStreamContext, accountID string) (nats.
}
// InitLeaderBucket creates (or attaches to) the shared spinifex-eks-leader
// bucket used for per-cluster reconciler leader-lease CAS locks. The bucket
// is configured with History=1 and a 60s TTL so stale leases expire on their
// own when a leader dies mid-cycle. utils.GetOrCreateKVBucket doesn't expose
// a TTL knob, so this function takes the direct js.CreateKeyValue path and
// falls back to js.KeyValue on already-exists.
func InitLeaderBucket(js nats.JetStreamContext) (nats.KeyValue, error) {
// bucket used for per-cluster reconciler leader-lease CAS locks, at the given
// replica count (clamped to a minimum of 1). The bucket is configured with
// History=1 and a 60s TTL so stale leases expire on their own when a leader
// dies mid-cycle. utils.GetOrCreateKVBucketWithReplicas doesn't expose a TTL
// knob, so this function sets Replicas directly on its own js.CreateKeyValue
// call and falls back to js.KeyValue on already-exists.
func InitLeaderBucket(js nats.JetStreamContext, replicas int) (nats.KeyValue, error) {
kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
Bucket: KVBucketEKSLeader,
History: 1,
TTL: KVBucketEKSLeaderTTL,
Bucket: KVBucketEKSLeader,
History: 1,
TTL: KVBucketEKSLeaderTTL,
Replicas: max(replicas, 1),
})
if err != nil {
kv, err = js.KeyValue(KVBucketEKSLeader)
+36 -4
View File
@@ -13,11 +13,11 @@ const testAccountID = "111122223333"
func TestGetOrCreateAccountBucket_Idempotent(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv1, err := GetOrCreateAccountBucket(js, testAccountID)
kv1, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
require.NotNil(t, kv1)
kv2, err := GetOrCreateAccountBucket(js, testAccountID)
kv2, err := GetOrCreateAccountBucket(js, testAccountID, 1)
require.NoError(t, err)
require.NotNil(t, kv2)
@@ -25,20 +25,52 @@ func TestGetOrCreateAccountBucket_Idempotent(t *testing.T) {
assert.Equal(t, kv1.Bucket(), kv2.Bucket())
}
// TestGetOrCreateAccountBucket_ReplicasClamped confirms the per-account
// bucket is created with the requested replica count, clamped to a minimum
// of 1. The embedded single-node test server rejects Replicas > 1
// ("replicas > 1 not supported in non-clustered mode"), so only the clamping
// path (replicas <= 0 -> 1) is exercisable end-to-end here; multi-node
// replica counts are exercised live (see the associated bug doc).
func TestGetOrCreateAccountBucket_ReplicasClamped(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv, err := GetOrCreateAccountBucket(js, testAccountID, 0)
require.NoError(t, err)
require.NotNil(t, kv)
si, err := js.StreamInfo("KV_" + AccountBucketName(testAccountID))
require.NoError(t, err)
assert.Equal(t, 1, si.Config.Replicas)
}
func TestInitLeaderBucket_Idempotent(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv1, err := InitLeaderBucket(js)
kv1, err := InitLeaderBucket(js, 1)
require.NoError(t, err)
require.NotNil(t, kv1)
assert.Equal(t, KVBucketEKSLeader, kv1.Bucket())
kv2, err := InitLeaderBucket(js)
kv2, err := InitLeaderBucket(js, 1)
require.NoError(t, err)
require.NotNil(t, kv2)
assert.Equal(t, KVBucketEKSLeader, kv2.Bucket())
}
// TestInitLeaderBucket_ReplicasClamped confirms the leader bucket is created
// with the requested replica count, clamped to a minimum of 1.
func TestInitLeaderBucket_ReplicasClamped(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv, err := InitLeaderBucket(js, -1)
require.NoError(t, err)
require.NotNil(t, kv)
si, err := js.StreamInfo("KV_" + KVBucketEKSLeader)
require.NoError(t, err)
assert.Equal(t, 1, si.Config.Replicas)
}
func TestKeyPaths_MatchQ2Spec(t *testing.T) {
cases := []struct {
name string
@@ -62,7 +62,7 @@ func newWebIdentityFixture(t *testing.T, svc *STSServiceImpl, accountID string)
federatedARN := handlers_iam.OIDCProviderARN(accountID, issuerHostPath)
// Publish JWKS to the per-cluster EKS bucket.
kv, err := handlers_eks.GetOrCreateAccountBucket(svc.js, accountID)
kv, err := handlers_eks.GetOrCreateAccountBucket(svc.js, accountID, 1)
require.NoError(t, err)
raw, err := json.Marshal(jwks)
require.NoError(t, err)
@@ -352,7 +352,7 @@ func TestAssumeRoleWithWebIdentity_ProviderNotRegistered(t *testing.T) {
federatedARN := handlers_iam.OIDCProviderARN(testCallerAccountID, issuerHostPath)
// Publish JWKS — but skip the IAM provider registration.
kv, err := handlers_eks.GetOrCreateAccountBucket(svc.js, testCallerAccountID)
kv, err := handlers_eks.GetOrCreateAccountBucket(svc.js, testCallerAccountID, 1)
require.NoError(t, err)
jwks := &JWKS{Keys: []JWK{{
Kty: "EC", Crv: "P-256", Alg: "ES256", Use: "sig", Kid: kid,
+2 -2
View File
@@ -79,7 +79,7 @@ func TestFetchClusterJWKS_NotFoundReturnsNil(t *testing.T) {
func TestFetchClusterJWKS_RoundTrip(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv, err := handlers_eks.GetOrCreateAccountBucket(js, "123456789012")
kv, err := handlers_eks.GetOrCreateAccountBucket(js, "123456789012", 1)
require.NoError(t, err)
want := &JWKS{Keys: []JWK{
@@ -102,7 +102,7 @@ func TestFetchClusterJWKS_RoundTrip(t *testing.T) {
func TestFetchClusterJWKS_EmptyKeysIsError(t *testing.T) {
_, _, js := testutil.StartTestJetStream(t)
kv, err := handlers_eks.GetOrCreateAccountBucket(js, "123456789012")
kv, err := handlers_eks.GetOrCreateAccountBucket(js, "123456789012", 1)
require.NoError(t, err)
_, err = kv.Put(handlers_eks.OIDCJWKSKey("empty-cluster"), []byte(`{"keys":[]}`))
+23 -11
View File
@@ -885,20 +885,32 @@ func launchService(cfg *Config) (err error) {
return
}
// Wait for 1 second to confirm nbdkit is running
time.Sleep(1 * time.Second)
// Poll for up to 1 second to confirm nbdkit is running, checking every
// 50ms so a fast failure is detected quickly instead of always
// costing the full second. Any exit within that window means NBDKit
// failed to stay up.
const nbdkitConfirmPollInterval = 50 * time.Millisecond
const nbdkitConfirmDeadline = 1 * time.Second
confirmDeadline := time.Now().Add(nbdkitConfirmDeadline)
for {
select {
case exitErr := <-exitChan:
ebsResponse.Error = fmt.Sprintf("nbdkit exited unexpectedly (code=%d)", exitErr)
respondAndPublish(msg, nc, "ebs.mount.response", ebsResponse)
return
default:
}
// Any exit within the first second means NBDKit failed to stay up.
select {
case exitErr := <-exitChan:
ebsResponse.Error = fmt.Sprintf("nbdkit exited unexpectedly (code=%d)", exitErr)
respondAndPublish(msg, nc, "ebs.mount.response", ebsResponse)
return
default:
// nbdkit is still running after 1 second, which means it started successfully
slog.InfoContext(ctx, "NBDKit started successfully and is running")
if time.Now().After(confirmDeadline) {
break
}
time.Sleep(nbdkitConfirmPollInterval)
}
// nbdkit is still running after the poll window, which means it started successfully
slog.InfoContext(ctx, "NBDKit started successfully and is running")
// NBDKit creates the socket with its own umask (typically 0755).
// The daemon (different user, same group) needs write access to connect.
if nbdSocket != "" {
+12 -3
View File
@@ -69,11 +69,20 @@ func IsAccountID(s string) bool {
return true
}
// GetOrCreateKVBucket creates or retrieves a NATS KV bucket.
// GetOrCreateKVBucket creates or retrieves a NATS KV bucket at Replicas=1.
func GetOrCreateKVBucket(js nats.JetStreamContext, bucketName string, history int) (nats.KeyValue, error) {
return GetOrCreateKVBucketWithReplicas(js, bucketName, history, 1)
}
// GetOrCreateKVBucketWithReplicas creates or retrieves a NATS KV bucket at the
// given replica count (clamped to a minimum of 1). The vendored nats.go
// (v1.52.0) has no UpdateKeyValue to live-upgrade an existing under-replicated
// bucket's replica count; existing buckets are upgraded on rebalance elsewhere.
func GetOrCreateKVBucketWithReplicas(js nats.JetStreamContext, bucketName string, history int, replicas int) (nats.KeyValue, error) {
kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
Bucket: bucketName,
History: SafeIntToUint8(history),
Bucket: bucketName,
History: SafeIntToUint8(history),
Replicas: max(replicas, 1),
})
if err != nil {
kv, err = js.KeyValue(bucketName)
+45
View File
@@ -106,6 +106,51 @@ func TestVersionStateMachine(t *testing.T) {
assert.Equal(t, "5", string(entry.Value()))
}
// streamReplicas returns the replica count of the JetStream stream backing a
// KV bucket, so tests can assert on the config actually sent to the server.
func streamReplicas(t *testing.T, js nats.JetStreamContext, bucket string) int {
t.Helper()
si, err := js.StreamInfo("KV_" + bucket)
require.NoError(t, err)
return si.Config.Replicas
}
func TestGetOrCreateKVBucket_CreatesAtReplicaOne(t *testing.T) {
_, js := startJSNATSServer(t)
kv, err := GetOrCreateKVBucket(js, "regression-bucket", 5)
require.NoError(t, err)
require.NotNil(t, kv)
assert.Equal(t, 1, streamReplicas(t, js, "regression-bucket"))
}
func TestGetOrCreateKVBucketWithReplicas_CreatesAtRequestedReplicas(t *testing.T) {
_, js := startJSNATSServer(t)
// The embedded single-node test server rejects Replicas > 1
// ("replicas > 1 not supported in non-clustered mode"), so the only
// replica count exercisable end-to-end here is 1; the clamping logic
// (max(replicas, 1)) is covered separately below.
kv, err := GetOrCreateKVBucketWithReplicas(js, "replicated-bucket", 5, 1)
require.NoError(t, err)
require.NotNil(t, kv)
assert.Equal(t, 1, streamReplicas(t, js, "replicated-bucket"))
}
func TestGetOrCreateKVBucketWithReplicas_ClampsBelowOne(t *testing.T) {
_, js := startJSNATSServer(t)
kv, err := GetOrCreateKVBucketWithReplicas(js, "clamped-zero", 1, 0)
require.NoError(t, err)
require.NotNil(t, kv)
assert.Equal(t, 1, streamReplicas(t, js, "clamped-zero"))
kv, err = GetOrCreateKVBucketWithReplicas(js, "clamped-negative", 1, -3)
require.NoError(t, err)
require.NotNil(t, kv)
assert.Equal(t, 1, streamReplicas(t, js, "clamped-negative"))
}
func TestDeleteKVBucketIfExists(t *testing.T) {
_, js := startJSNATSServer(t)
+28 -21
View File
@@ -72,7 +72,22 @@ func ForceKillProcess(pid int, timeout time.Duration) error {
return WaitForProcessExit(pid, timeout)
}
// killProcessPollInterval and killProcessGracePeriod drive KillProcess's
// liveness poll: check every pollInterval instead of blocking for the full
// gracePeriod, so a process that dies quickly is detected almost immediately.
const (
killProcessPollInterval = 50 * time.Millisecond
killProcessGracePeriod = 120 * time.Second
)
func KillProcess(pid int) error {
return killProcessWithTiming(pid, killProcessPollInterval, killProcessGracePeriod)
}
// killProcessWithTiming implements KillProcess with an injectable poll
// interval and grace period so tests can exercise the SIGKILL escalation
// path without waiting out the full production grace period.
func killProcessWithTiming(pid int, pollInterval, gracePeriod time.Duration) error {
process, err := os.FindProcess(pid)
if err != nil {
@@ -84,33 +99,25 @@ func KillProcess(pid int) error {
return err
}
checks := 0
deadline := time.Now().Add(gracePeriod)
for {
time.Sleep(1 * time.Second)
process, err = os.FindProcess(pid)
if err != nil {
return err
if !ProcessAlive(pid) {
return nil // process has terminated
}
err = process.Signal(syscall.Signal(0))
if err != nil {
break // Signal(0) error means process has terminated
}
checks++
if checks > 120 {
err = process.Kill() // SIGKILL after 120 s
if time.Now().After(deadline) {
process, err = os.FindProcess(pid)
if err != nil {
return err
}
break
}
}
if err := process.Kill(); err != nil { // SIGKILL after grace period
return err
}
return nil //nolint:nilerr // Signal(0) error means process exited — that's success
return nil
}
time.Sleep(pollInterval)
}
}
+78
View File
@@ -0,0 +1,78 @@
package utils
import (
"bufio"
"os/exec"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
)
// KillProcess must poll for liveness instead of blindly sleeping, so a
// process that dies quickly (nbdkit on unmount, in production) is reaped
// in milliseconds rather than costing a fixed 1s+ per call.
func TestKillProcessReturnsQuicklyWhenProcessExitsPromptly(t *testing.T) {
cmd := exec.Command("sleep", "60") // dies immediately on default SIGTERM
require.NoError(t, cmd.Start())
pid := cmd.Process.Pid
done := make(chan struct{})
go func() { _ = cmd.Wait(); close(done) }()
start := time.Now()
err := KillProcess(pid)
elapsed := time.Since(start)
require.NoError(t, err)
assert.Less(t, elapsed, 300*time.Millisecond,
"a process that dies promptly must be reaped well under the old fixed 1s sleep")
assert.False(t, ProcessAlive(pid))
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("process did not exit after KillProcess returned")
}
}
// The escalation-to-SIGKILL path must still be bounded, not hang forever,
// when a process ignores SIGTERM. Uses a shortened grace period so the test
// doesn't have to wait out the real (much longer) production grace period.
func TestKillProcessEscalatesToSigkillWithinDeadline(t *testing.T) {
// Ignore SIGTERM so the process only dies via the SIGKILL escalation.
// The loop (rather than a single "sleep") stops the shell from exec-ing
// straight into sleep, which would silently drop the trap. Wait for the
// "ready" line so we never signal before the trap is installed.
cmd := exec.Command("sh", "-c", "trap '' TERM; echo ready; while true; do sleep 1; done")
stdout, err := cmd.StdoutPipe()
require.NoError(t, err)
require.NoError(t, cmd.Start())
pid := cmd.Process.Pid
done := make(chan struct{})
go func() { _ = cmd.Wait(); close(done) }()
scanner := bufio.NewScanner(stdout)
require.True(t, scanner.Scan(), "expected ready line from test process")
require.Equal(t, "ready", scanner.Text())
const pollInterval = 20 * time.Millisecond
const gracePeriod = 300 * time.Millisecond
start := time.Now()
err = killProcessWithTiming(pid, pollInterval, gracePeriod)
elapsed := time.Since(start)
require.NoError(t, err)
assert.GreaterOrEqual(t, elapsed, gracePeriod,
"must wait out the full grace period before escalating")
assert.Less(t, elapsed, gracePeriod+2*time.Second,
"escalation must be bounded near the grace period, not hang indefinitely")
select {
case <-done:
case <-time.After(2 * time.Second):
t.Fatal("process did not exit after KillProcess escalated to SIGKILL")
}
assert.False(t, ProcessAlive(pid), "process must be gone after SIGKILL escalation")
}
+7
View File
@@ -46,3 +46,10 @@ var ErrENIPipelineTimeout = errors.New("ENI hot-plug pipeline timed out")
// instance has no live QMPClient (daemon restart mid-launch, terminated
// VM, etc).
var ErrQMPUnavailable = errors.New("QMP unavailable for instance")
// ErrStoppedInstanceClaimed is returned by StateStore.ClaimStoppedInstance
// when a concurrent caller already claimed (or the record was otherwise
// removed from) the stopped-instance store for the given id. The caller
// lost the race and must not proceed to allocate resources or launch qemu
// for that instance.
var ErrStoppedInstanceClaimed = errors.New("stopped instance already claimed")
+12 -9
View File
@@ -197,8 +197,8 @@ func (m *Manager) launch(ctx context.Context, instance *VM) (err error) {
}
}
// Mark boot volumes as "in-use" now that instance is confirmed running.
m.markBootVolumesInUse(instance)
// Mark attached volumes as "in-use" now that instance is confirmed running.
m.markAttachedVolumesInUse(instance)
if m.deps.Hooks.OnInstanceUp != nil {
// Launch path: per-instance subscribe failures are logged and the
@@ -214,21 +214,24 @@ func (m *Manager) launch(ctx context.Context, instance *VM) (err error) {
return nil
}
// markBootVolumesInUse re-asserts "in-use" status for an instance's boot
// volumes once it is confirmed running. Used by the launch path and the
// daemon-restart reconnect path so both keep volume state consistent with a
// running instance. Errors are logged, not fatal.
func (m *Manager) markBootVolumesInUse(instance *VM) {
// markAttachedVolumesInUse re-asserts "in-use" status for every volume an
// instance currently has attached (boot and non-boot) once it is confirmed
// running. Used by the launch path and the daemon-restart reconnect path so
// both keep volume state consistent with a running instance — otherwise a
// non-boot volume (e.g. an EKS stateful-pod data volume) can read "available"
// while the instance still has it attached. Errors are logged, not fatal.
func (m *Manager) markAttachedVolumesInUse(instance *VM) {
if m.deps.VolumeStateUpdater == nil {
return
}
instance.EBSRequests.Mu.Lock()
defer instance.EBSRequests.Mu.Unlock()
for _, ebsReq := range instance.EBSRequests.Requests {
if !ebsReq.Boot {
// EFI pflash is not a KV-tracked EBS volume; skip it like the unmount gate.
if ebsReq.EFI {
continue
}
if err := m.deps.VolumeStateUpdater.UpdateVolumeState(ebsReq.Name, "in-use", instance.ID, ""); err != nil {
if err := m.deps.VolumeStateUpdater.UpdateVolumeState(ebsReq.Name, "in-use", instance.ID, ebsReq.DeviceName); err != nil {
slog.Error("Failed to update volume state to in-use", "volumeId", ebsReq.Name, "err", err)
}
}
+11 -2
View File
@@ -55,10 +55,19 @@ type Deps struct {
BindHost string
// DetachDelay is the pause between QMP device_del and blockdev-del
// during DetachVolume, and the polling interval for blockdev-del retry.
// Zero is acceptable in tests; production uses 1s.
// during DetachVolume when DeviceDeletedTimeout is unset, and the
// polling interval for blockdev-del retry. Zero is acceptable in tests;
// production uses 1s.
DetachDelay time.Duration
// DeviceDeletedTimeout bounds how long DetachVolume waits for QEMU's
// DEVICE_DELETED event after device_del before falling back to the
// blockdev-del retry loop. This is the real signal that the guest has
// released the block node; when set (>0) it replaces the DetachDelay
// sleep for the common path. Zero disables the wait (old sleep+retry
// behavior); production uses 15s.
DeviceDeletedTimeout time.Duration
// ConsumeCleanShutdownMarker reports whether the previous run wrote a clean
// shutdown marker, consuming it. Nil is treated as "no marker" (cautious recovery).
ConsumeCleanShutdownMarker func() bool
+29
View File
@@ -1,6 +1,7 @@
package vm
import (
"errors"
"maps"
"sync"
@@ -81,6 +82,20 @@ func (f *fakeStateStore) ListStoppedInstances() ([]*VM, error) {
return out, nil
}
// ClaimStoppedInstance mimics the real atomic-delete claim for tests: under
// the store lock, remove and return the entry, or ErrStoppedInstanceClaimed
// if it is already gone (claimed by a concurrent caller, or never existed).
func (f *fakeStateStore) ClaimStoppedInstance(id string) (*VM, error) {
f.mu.Lock()
defer f.mu.Unlock()
v, ok := f.stopped[id]
if !ok {
return nil, ErrStoppedInstanceClaimed
}
delete(f.stopped, id)
return v, nil
}
func (f *fakeStateStore) WriteTerminatedInstance(id string, v *VM) error {
f.mu.Lock()
defer f.mu.Unlock()
@@ -105,6 +120,20 @@ func (f *fakeStateStore) DeleteTerminatedInstance(id string) error {
return nil
}
// UpdateTerminatedInstance mimics the real CAS semantics for tests: mutate
// runs under the store lock against the stored value, matching the
// read-modify-write contract without needing a real KV revision.
func (f *fakeStateStore) UpdateTerminatedInstance(id string, mutate func(*VM)) (*VM, error) {
f.mu.Lock()
defer f.mu.Unlock()
v, ok := f.terminated[id]
if !ok {
return nil, errors.New("terminated instance not found")
}
mutate(v)
return v, nil
}
var _ StateStore = (*fakeStateStore)(nil)
// fakeVolumeMounter records every call so lifecycle tests can assert ordering.
+4 -3
View File
@@ -320,9 +320,10 @@ func (m *Manager) reconnectInstance(instance *VM) error {
instance.Status = StateRunning
// Re-assert boot-volume in-use state; a daemon restart otherwise leaves the
// root volume "available" while the instance runs (split-brain, siv-464).
m.markBootVolumesInUse(instance)
// Re-assert in-use state for every attached volume; a daemon restart
// otherwise leaves boot and non-boot volumes "available" while the
// instance runs (split-brain).
m.markAttachedVolumesInUse(instance)
if err := m.writeRunningState(); err != nil {
return fmt.Errorf("failed to persist reconnected instance state: %w", err)
+7 -5
View File
@@ -885,7 +885,7 @@ func TestReconnectInstance(t *testing.T) {
assert.NotNil(t, saved["i-no-hook"], "the reconnected instance must appear in the persisted snapshot")
})
t.Run("reconnect re-asserts boot-volume in-use, skips non-boot", func(t *testing.T) {
t.Run("reconnect re-asserts in-use for boot and non-boot volumes", func(t *testing.T) {
stateUpdater := &fakeVolumeStateUpdater{}
m := NewManager()
m.SetDeps(Deps{
@@ -901,17 +901,19 @@ func TestReconnectInstance(t *testing.T) {
instance := &VM{ID: "i-reconnect-vol", Status: StatePending}
instance.EBSRequests.Requests = []types.EBSRequest{
{Name: "vol-root", Boot: true},
{Name: "vol-data", Boot: false},
{Name: "vol-data", Boot: false, DeviceName: "/dev/sdf"},
}
m.Insert(instance)
require.NoError(t, m.reconnectInstance(instance))
calls := stateUpdater.snapshot()
require.Len(t, calls, 1,
"reconnect must re-assert exactly the boot volume — a daemon restart otherwise leaves the root volume 'available' while the VM runs (siv-464)")
require.Len(t, calls, 2,
"reconnect must re-assert every attached volume — a daemon restart otherwise leaves attached volumes 'available' while the VM runs, wedging stateful workloads that check volume availability")
assert.Equal(t, volumeStateUpdate{VolumeID: "vol-root", State: "in-use", InstanceID: "i-reconnect-vol"}, calls[0],
"the boot volume must be marked in-use under the running instance; non-boot volumes are owned by their own attach flow")
"the boot volume must be marked in-use under the running instance with no API device (boot volumes carry none)")
assert.Equal(t, volumeStateUpdate{VolumeID: "vol-data", State: "in-use", InstanceID: "i-reconnect-vol", AttachmentDevice: "/dev/sdf"}, calls[1],
"non-boot volumes must also be marked in-use, carrying the persisted API-form device name (not a guest/nbd path)")
})
t.Run("AttachQMP succeeds, OnInstanceUp errors: QMP closed and nil'd", func(t *testing.T) {
+10 -4
View File
@@ -1172,13 +1172,19 @@ func (s *signalingStore) SaveRunningState(string, map[string]*VM) error {
func (s *signalingStore) LoadRunningState(string) (map[string]*VM, error) {
return map[string]*VM{}, nil
}
func (s *signalingStore) WriteStoppedInstance(string, *VM) error { return nil }
func (s *signalingStore) LoadStoppedInstance(string) (*VM, error) { return nil, nil }
func (s *signalingStore) DeleteStoppedInstance(string) error { return nil }
func (s *signalingStore) ListStoppedInstances() ([]*VM, error) { return nil, nil }
func (s *signalingStore) WriteStoppedInstance(string, *VM) error { return nil }
func (s *signalingStore) LoadStoppedInstance(string) (*VM, error) { return nil, nil }
func (s *signalingStore) DeleteStoppedInstance(string) error { return nil }
func (s *signalingStore) ListStoppedInstances() ([]*VM, error) { return nil, nil }
func (s *signalingStore) ClaimStoppedInstance(string) (*VM, error) {
return nil, ErrStoppedInstanceClaimed
}
func (s *signalingStore) WriteTerminatedInstance(string, *VM) error { return nil }
func (s *signalingStore) ListTerminatedInstances() ([]*VM, error) { return nil, nil }
func (s *signalingStore) DeleteTerminatedInstance(string) error { return nil }
func (s *signalingStore) UpdateTerminatedInstance(string, func(*VM)) (*VM, error) {
return nil, nil
}
var _ StateStore = (*signalingStore)(nil)
+12
View File
@@ -16,8 +16,20 @@ type StateStore interface {
LoadStoppedInstance(id string) (*VM, error)
DeleteStoppedInstance(id string) error
ListStoppedInstances() ([]*VM, error)
// ClaimStoppedInstance atomically removes id's record from the shared
// store and returns the VM it held, so that at most one caller across
// the cluster can ever win a race to (re)launch the same stopped
// instance. Returns ErrStoppedInstanceClaimed if a concurrent caller
// already claimed (or otherwise removed) the record.
ClaimStoppedInstance(id string) (*VM, error)
WriteTerminatedInstance(id string, v *VM) error
ListTerminatedInstances() ([]*VM, error)
DeleteTerminatedInstance(id string) error
// UpdateTerminatedInstance atomically applies mutate to the current
// record for id and writes it back using optimistic concurrency (CAS),
// so a concurrent writer to the same record (e.g. two teardown-reaper
// sweeps advancing different dependents) cannot clobber the other's
// progress. Returns an error if no record exists for id.
UpdateTerminatedInstance(id string, mutate func(*VM)) (*VM, error)
}
+8 -3
View File
@@ -38,11 +38,16 @@ func (m *Manager) markTeardown(instance *VM, dep string, state TeardownState) {
// markTeardownResult records done on nil error, failed otherwise.
func (m *Manager) markTeardownResult(instance *VM, dep string, err error) {
m.markTeardown(instance, dep, resultState(err))
}
// resultState maps a cleaner call's error into its teardown mark: done on
// nil, failed otherwise.
func resultState(err error) TeardownState {
if err != nil {
m.markTeardown(instance, dep, TeardownFailed)
return
return TeardownFailed
}
m.markTeardown(instance, dep, TeardownDone)
return TeardownDone
}
// TeardownComplete reports whether every tracked dependent is done. A terminated
+40 -17
View File
@@ -59,7 +59,13 @@ func (r *TerminatedTeardownReaper) Sweep(context.Context) (int, error) {
continue // already done on arrival: leave to the 1h bucket TTL
}
r.retryOutstanding(v)
results := r.retryOutstanding(v)
for dep, state := range results {
if v.Teardown == nil {
v.Teardown = make(map[string]string)
}
v.Teardown[dep] = string(state)
}
if v.TeardownComplete() && time.Since(v.TerminatedAt) >= terminatedVisibilityWindow {
if r.purge(v) {
@@ -68,40 +74,55 @@ func (r *TerminatedTeardownReaper) Sweep(context.Context) (int, error) {
continue
}
// Either still incomplete (retry next sweep), or complete but inside the
// describe-visibility window: persist progress and leave the record to
// the 1h bucket TTL so the instance stays describable as terminated.
if err := r.m.deps.StateStore.WriteTerminatedInstance(v.ID, v); err != nil {
slog.Warn("vm/gc: failed to persist advanced teardown marks",
"instanceId", v.ID, "err", err)
// describe-visibility window: merge this sweep's marks into whatever is
// currently in KV via CAS, so a concurrent writer advancing a different
// dependent for the same record isn't clobbered by our local snapshot.
if len(results) > 0 {
if _, err := r.m.deps.StateStore.UpdateTerminatedInstance(v.ID, func(fresh *VM) {
if fresh.Teardown == nil {
fresh.Teardown = make(map[string]string)
}
for dep, state := range results {
fresh.Teardown[dep] = string(state)
}
}); err != nil {
slog.Warn("vm/gc: failed to persist advanced teardown marks",
"instanceId", v.ID, "err", err)
}
}
}
return reaped, nil
}
// retryOutstanding re-drives every not-`done` dependent through the idempotent
// cleaner and re-stamps the result. The cleaner calls are idempotent (absent →
// success), so a successful re-drive confirms the Spinifex-side teardown; the
// dataplane object (OVN LSP, NAT rule) is pruned by the cluster reconciler.
func (r *TerminatedTeardownReaper) retryOutstanding(v *VM) {
// cleaner and returns the dep→result map for whatever was retried. The cleaner
// calls are idempotent (absent → success), so a successful re-drive confirms
// the Spinifex-side teardown; the dataplane object (OVN LSP, NAT rule) is
// pruned by the cluster reconciler. Pure w.r.t. KV/local state — callers merge
// the results atomically via StateStore.UpdateTerminatedInstance so the cleaner
// calls (which have real side effects) never run twice for a single CAS retry.
func (r *TerminatedTeardownReaper) retryOutstanding(v *VM) map[string]TeardownState {
c := r.m.deps.InstanceCleaner
if c == nil {
return
return nil
}
results := make(map[string]TeardownState)
if outstanding(v, TeardownVolumes) {
r.m.markTeardownResult(v, TeardownVolumes, c.DeleteVolumes(v))
results[TeardownVolumes] = resultState(c.DeleteVolumes(v))
}
if outstanding(v, TeardownGPU) {
r.m.markTeardownResult(v, TeardownGPU, c.ReleaseGPU(v))
results[TeardownGPU] = resultState(c.ReleaseGPU(v))
}
if outstanding(v, TeardownPlacement) {
r.m.markTeardownResult(v, TeardownPlacement, c.RemoveFromPlacementGroup(v))
results[TeardownPlacement] = resultState(c.RemoveFromPlacementGroup(v))
}
// NAT: re-publishes vpc.delete-nat + frees the IPAM slot. The cluster
// reconciler heals the dataplane NAT rule.
if outstanding(v, TeardownNAT) {
r.m.markTeardownResult(v, TeardownNAT, c.ReleasePublicIP(v))
results[TeardownNAT] = resultState(c.ReleasePublicIP(v))
}
// ENI delete + OVN: deleting the ENI KV record turns its LSP into an orphan
@@ -109,9 +130,11 @@ func (r *TerminatedTeardownReaper) retryOutstanding(v *VM) {
// both eni and ovn.
if outstanding(v, TeardownENI) || outstanding(v, TeardownOVN) {
eniErr := c.DetachAndDeleteENI(v)
r.m.markTeardownResult(v, TeardownENI, eniErr)
r.m.markTeardownResult(v, TeardownOVN, eniErr)
results[TeardownENI] = resultState(eniErr)
results[TeardownOVN] = resultState(eniErr)
}
return results
}
func (r *TerminatedTeardownReaper) purge(v *VM) bool {
+74 -4
View File
@@ -15,9 +15,10 @@ import (
"github.com/mulgadc/spinifex/spinifex/utils"
)
// blockdevDelMaxAttempts caps the bounded retry on "node is in use" while
// QEMU drains pending I/O after device_del. 20 × DetachDelay (default 1s)
// gives a 20s budget.
// blockdevDelMaxAttempts caps the bounded retry on "node is in use" as a
// safety net after the DEVICE_DELETED wait (see DeviceDeletedTimeout): a miss
// on the event (timeout, or a client not opted in) falls back to polling at
// DetachDelay. 20 × DetachDelay (default 1s) gives a 20s budget.
const blockdevDelMaxAttempts = 20
// rollbackUnmount best-effort unmounts a volume while unwinding a failed
@@ -320,12 +321,14 @@ func (m *Manager) DetachVolume(ctx context.Context, id, volumeID, device string,
// device_del is idempotent on DeviceNotFound so a second AWS-CLI
// retry can drive blockdev-del to completion when a prior detach left
// the guest device gone but the block node intact.
deviceDelIssued := false
_, err := sendQMPCommand(ctx, instance.QMPClient, qmp.QMPCommand{
Execute: "device_del",
Arguments: map[string]any{"id": deviceID},
}, instance.ID)
switch {
case err == nil:
deviceDelIssued = true
case isQMPDeviceNotFound(err):
slog.InfoContext(ctx, "DetachVolume: guest device already removed (resuming detach)",
"volumeId", volumeID, "err", err)
@@ -337,7 +340,18 @@ func (m *Manager) DetachVolume(ctx context.Context, id, volumeID, device string,
return "", fmt.Errorf("QMP device_del: %w", err)
}
if m.deps.DetachDelay > 0 {
// device_del only requests the unplug; QEMU frees the block node once the
// guest ACKs it (DEVICE_DELETED). Wait for that real completion signal
// instead of blindly sleeping, so blockdev-del isn't attempted while the
// node still has a user. A miss (timeout, event drained by a racing read,
// or a resumed/forced detach with no fresh unplug in flight) falls
// through to the bounded retry below unchanged.
if deviceDelIssued && m.deps.DeviceDeletedTimeout > 0 {
if waitErr := waitForDeviceDeletedEvent(ctx, instance.QMPClient, deviceID, m.deps.DeviceDeletedTimeout, instance.ID); waitErr != nil {
slog.DebugContext(ctx, "DetachVolume: DEVICE_DELETED not observed, falling back to blockdev-del retry",
"volumeId", volumeID, "err", waitErr)
}
} else if m.deps.DetachDelay > 0 {
time.Sleep(m.deps.DetachDelay)
}
@@ -408,6 +422,62 @@ func (m *Manager) DetachVolume(ctx context.Context, id, volumeID, device string,
return ebsReq.DeviceName, nil
}
// waitForDeviceDeletedEvent blocks until QEMU emits a DEVICE_DELETED event
// matching deviceID or timeout elapses, returning nil only on a match. It
// takes over the QMP connection directly (no in-flight command is expected
// while it runs) rather than going through sendQMPCommand, since the signal
// we need is an async event, not a command reply.
//
// A timeout is treated the same as any other decode failure: q.Decoder is a
// single shared *json.Decoder, and once Decode returns an error — including a
// read-deadline timeout — that Decoder instance is permanently unusable, so
// q.Dead is set exactly as sendQMPCommand does. The next QMP command
// transparently redials via reconnectQMP. This is non-fatal either way:
// DetachVolume falls back to the bounded blockdev-del retry, so a missed
// event never blocks the caller forever.
func waitForDeviceDeletedEvent(ctx context.Context, q *qmp.QMPClient, deviceID string, timeout time.Duration, instanceID string) error {
if q == nil || q.Conn == nil || q.Decoder == nil {
return fmt.Errorf("QMP client is not initialized")
}
q.Mu.Lock()
defer q.Mu.Unlock()
if q.Dead {
return fmt.Errorf("QMP client is disconnected")
}
if err := q.Conn.SetReadDeadline(time.Now().Add(timeout)); err != nil {
return fmt.Errorf("set read deadline: %w", err)
}
defer func() { _ = q.Conn.SetReadDeadline(time.Time{}) }()
for {
var msg struct {
Event string `json:"event"`
Data struct {
Device string `json:"device"`
Path string `json:"path"`
} `json:"data"`
}
if err := q.Decoder.Decode(&msg); err != nil {
q.Dead = true
return fmt.Errorf("decode error waiting for DEVICE_DELETED: %w", err)
}
if msg.Event != "DEVICE_DELETED" {
// Nothing else should be in flight on this connection while we
// hold q.Mu, but tolerate and skip any stray message rather than
// treating it as our signal.
continue
}
if msg.Data.Device == deviceID || strings.Contains(msg.Data.Path, deviceID) {
return nil
}
slog.DebugContext(ctx, "waitForDeviceDeletedEvent: DEVICE_DELETED for a different device, still waiting",
"wantDevice", deviceID, "gotDevice", msg.Data.Device, "instanceId", instanceID)
}
}
// tryBlockdevDel issues blockdev-del with bounded retry on "is in use" errors.
// QEMU surfaces this as GenericError while NBD teardown and in-flight I/O drain.
func (m *Manager) tryBlockdevDel(ctx context.Context, instance *VM, nodeName string) error {
+350
View File
@@ -7,9 +7,11 @@ import (
"fmt"
"log/slog"
"net"
"path/filepath"
"sync"
"sync/atomic"
"testing"
"time"
"github.com/aws/aws-sdk-go/aws"
"github.com/aws/aws-sdk-go/service/ec2"
@@ -64,6 +66,119 @@ func newMockQMPClient(t *testing.T, responder func(qmp.QMPCommand) map[string]an
return client, cancel
}
// mockQMPServer lets a responder push asynchronous QMP events (e.g.
// DEVICE_DELETED) onto the same stream as command replies. Writes are
// serialized so a delayed event push never interleaves with a reply.
type mockQMPServer struct {
mu sync.Mutex
enc *json.Encoder
}
func (s *mockQMPServer) pushEvent(ev map[string]any) error {
s.mu.Lock()
defer s.mu.Unlock()
return s.enc.Encode(ev)
}
// newMockQMPClientWithEvents is newMockQMPClient plus a *mockQMPServer handle
// so the responder can schedule out-of-band events (typically from a
// goroutine) in addition to its synchronous command reply.
//
// This uses a real unix socket rather than net.Pipe: net.Pipe's deadline
// implementation does not reliably support "timeout, then clear the deadline,
// then read again" on the same conn (verified against Go's net.Pipe — a real
// unix socket, matching production QMP transport, does not have this
// limitation), and DEVICE_DELETED-wait tests need exactly that sequence.
//
// The listener keeps accepting connections for the life of the test (each
// served on its own goroutine, greeting first) so a client-side reconnect —
// e.g. reconnectQMP redialing after waitForDeviceDeletedEvent marks the
// client Dead on timeout — is served exactly like the first connection.
func newMockQMPClientWithEvents(t *testing.T, responder func(qmp.QMPCommand, *mockQMPServer) map[string]any) (*qmp.QMPClient, func()) {
t.Helper()
sockPath := filepath.Join(t.TempDir(), "qmp.sock")
ln, err := net.Listen("unix", sockPath)
require.NoError(t, err)
var wg sync.WaitGroup
serveConn := func(conn net.Conn) {
defer wg.Done()
defer conn.Close()
srv := &mockQMPServer{enc: json.NewEncoder(conn)}
if err := srv.pushEvent(map[string]any{
"QMP": map[string]any{
"version": map[string]any{"qemu": map[string]any{"major": 8, "minor": 0}},
"capabilities": []string{},
},
}); err != nil {
return
}
dec := json.NewDecoder(conn)
for {
var cmd qmp.QMPCommand
if err := dec.Decode(&cmd); err != nil {
return
}
if cmd.Execute == "qmp_capabilities" {
if err := srv.pushEvent(map[string]any{"return": map[string]any{}}); err != nil {
return
}
continue
}
var reply map[string]any
if responder != nil {
reply = responder(cmd, srv)
}
if reply == nil {
reply = map[string]any{"return": map[string]any{}}
}
if err := srv.pushEvent(reply); err != nil {
return
}
}
}
acceptDone := make(chan struct{})
go func() {
defer close(acceptDone)
for {
conn, err := ln.Accept()
if err != nil {
return
}
wg.Add(1)
go serveConn(conn)
}
}()
clientConn, err := net.Dial("unix", sockPath)
require.NoError(t, err)
dec := json.NewDecoder(clientConn)
var greeting qmp.QMPGreeting
require.NoError(t, dec.Decode(&greeting), "mock QMP server greeting")
client := &qmp.QMPClient{
Conn: clientConn,
Decoder: dec,
Encoder: json.NewEncoder(clientConn),
Path: sockPath,
}
cancel := func() {
_ = ln.Close()
// client.Conn may have been swapped by reconnectQMP after a Dead
// redial, so close the current field value, not the initial dial.
_ = client.Conn.Close()
<-acceptDone
wg.Wait()
}
return client, cancel
}
// qmpRecorder records the order of QMP commands for assertion in rollback
// tests. Safe for concurrent use because sendQMPCommand serializes via the
// QMPClient mutex, but we lock anyway for clarity.
@@ -911,3 +1026,238 @@ func TestTryBlockdevDel(t *testing.T) {
"retry must cap at blockdevDelMaxAttempts attempts")
})
}
// newMockQMPServerConn wires a QMPClient over net.Pipe with no command-reply
// loop at all — the test drives the server side directly via the returned
// json.Encoder, which is what waitForDeviceDeletedEvent tests need: an event
// pushed independently of any command/response exchange.
func newMockQMPServerConn(t *testing.T) (*qmp.QMPClient, *json.Encoder, func()) {
t.Helper()
sockPath := filepath.Join(t.TempDir(), "qmp.sock")
ln, err := net.Listen("unix", sockPath)
require.NoError(t, err)
accepted := make(chan net.Conn, 1)
go func() {
conn, err := ln.Accept()
if err != nil {
close(accepted)
return
}
accepted <- conn
}()
clientConn, err := net.Dial("unix", sockPath)
require.NoError(t, err)
serverConn := <-accepted
_ = ln.Close()
require.NotNil(t, serverConn, "mock QMP server failed to accept")
client := &qmp.QMPClient{
Conn: clientConn,
Decoder: json.NewDecoder(clientConn),
Encoder: json.NewEncoder(clientConn),
Path: sockPath,
}
enc := json.NewEncoder(serverConn)
cancel := func() {
_ = clientConn.Close()
_ = serverConn.Close()
}
return client, enc, cancel
}
// TestWaitForDeviceDeletedEvent covers waitForDeviceDeletedEvent directly:
// matching event, non-matching event followed by a match, and timeout. Uses
// a short timeout so the timeout sub-test stays fast.
func TestWaitForDeviceDeletedEvent(t *testing.T) {
t.Run("matching event returns nil", func(t *testing.T) {
qmpClient, enc, cancel := newMockQMPServerConn(t)
defer cancel()
errCh := make(chan error, 1)
go func() {
errCh <- waitForDeviceDeletedEvent(t.Context(), qmpClient, "vdisk-vol-1", 2*time.Second, "i-1")
}()
require.NoError(t, enc.Encode(map[string]any{
"event": "DEVICE_DELETED",
"data": map[string]any{"device": "vdisk-vol-1", "path": "/machine/peripheral/vdisk-vol-1"},
}))
require.NoError(t, <-errCh)
})
t.Run("non-matching event is skipped, later match returns nil", func(t *testing.T) {
qmpClient, enc, cancel := newMockQMPServerConn(t)
defer cancel()
errCh := make(chan error, 1)
go func() {
errCh <- waitForDeviceDeletedEvent(t.Context(), qmpClient, "vdisk-vol-2", 2*time.Second, "i-1")
}()
require.NoError(t, enc.Encode(map[string]any{
"event": "DEVICE_DELETED",
"data": map[string]any{"device": "vdisk-vol-1", "path": "/machine/peripheral/vdisk-vol-1"},
}))
require.NoError(t, enc.Encode(map[string]any{
"event": "DEVICE_DELETED",
"data": map[string]any{"device": "vdisk-vol-2", "path": "/machine/peripheral/vdisk-vol-2"},
}))
require.NoError(t, <-errCh)
})
t.Run("no event before timeout returns error and marks the client dead", func(t *testing.T) {
qmpClient, _, cancel := newMockQMPServerConn(t)
defer cancel()
err := waitForDeviceDeletedEvent(t.Context(), qmpClient, "vdisk-vol-1", 50*time.Millisecond, "i-1")
require.Error(t, err)
// q.Decoder is a single shared *json.Decoder: once Decode returns any
// error (including a deadline timeout) that instance is permanently
// unusable, so Dead must be set — exactly like sendQMPCommand — so the
// next QMP command redials via reconnectQMP instead of reusing it.
assert.True(t, qmpClient.Dead,
"a decode timeout must mark the client dead so the next command reconnects")
})
t.Run("nil client returns error", func(t *testing.T) {
err := waitForDeviceDeletedEvent(t.Context(), nil, "vdisk-vol-1", time.Second, "i-1")
require.Error(t, err)
})
}
// TestDetachVolume_WaitsForDeviceDeletedBeforeBlockdevDel is the regression
// guard for the "node in use" wedge: blockdev-del only ever succeeds once
// DEVICE_DELETED has been observed, so if DetachVolume attempted it before
// waiting for the event, this test would exhaust blockdevDelMaxAttempts and
// fail exactly like the reported bug. Release of the event is gated on an
// explicit signal (not a race against wall-clock timing) so the test is
// deterministic; a minimum-elapsed-time assertion confirms DetachVolume
// actually blocked on the event rather than skipping the wait.
func TestDetachVolume_WaitsForDeviceDeletedBeforeBlockdevDel(t *testing.T) {
var blockdevAttempts atomic.Int32
releaseEvent := make(chan struct{})
qmpClient, cancel := newMockQMPClientWithEvents(t, func(cmd qmp.QMPCommand, srv *mockQMPServer) map[string]any {
switch cmd.Execute {
case "device_del":
go func() {
<-releaseEvent
_ = srv.pushEvent(map[string]any{
"event": "DEVICE_DELETED",
"data": map[string]any{"device": "vdisk-vol-1", "path": "/machine/peripheral/vdisk-vol-1"},
})
}()
return nil
case "blockdev-del":
blockdevAttempts.Add(1)
return nil
}
return nil
})
defer cancel()
mounter := &fakeVolumeMounter{}
m := NewManagerWithDeps(Deps{VolumeMounter: mounter, DeviceDeletedTimeout: 5 * time.Second})
m.Insert(&VM{
ID: "i-1",
Status: StateRunning,
Instance: &ec2.Instance{},
QMPClient: qmpClient,
EBSRequests: types.EBSRequests{
Requests: []types.EBSRequest{{Name: "vol-1", DeviceName: "/dev/sdf"}},
},
})
const releaseDelay = 100 * time.Millisecond
go func() {
time.Sleep(releaseDelay)
close(releaseEvent)
}()
start := time.Now()
device, err := m.DetachVolume(t.Context(), "i-1", "vol-1", "", false)
elapsed := time.Since(start)
require.NoError(t, err)
assert.Equal(t, "/dev/sdf", device)
assert.Equal(t, int32(1), blockdevAttempts.Load(),
"waiting for DEVICE_DELETED before the first blockdev-del attempt must avoid any 'node in use' retry")
assert.GreaterOrEqual(t, elapsed, releaseDelay/2,
"DetachVolume must actually block on the DEVICE_DELETED event rather than proceeding immediately")
}
// TestDetachVolume_DeviceDeletedTimeout_FallsBackToRetry covers the case
// where DEVICE_DELETED is never observed within DeviceDeletedTimeout (e.g. a
// lost event, or a QEMU version quirk): DetachVolume must still fall back to
// the bounded blockdev-del retry and succeed once the node frees up, rather
// than surfacing the wait timeout as a hard failure.
func TestDetachVolume_DeviceDeletedTimeout_FallsBackToRetry(t *testing.T) {
var blockdevAttempts atomic.Int32
qmpClient, cancel := newMockQMPClientWithEvents(t, func(cmd qmp.QMPCommand, srv *mockQMPServer) map[string]any {
if cmd.Execute == "blockdev-del" {
n := blockdevAttempts.Add(1)
if n < 2 {
return map[string]any{"error": map[string]any{"class": "GenericError", "desc": "Node 'nbd-vol-1' is in use"}}
}
}
return nil
})
defer cancel()
mounter := &fakeVolumeMounter{}
m := NewManagerWithDeps(Deps{VolumeMounter: mounter, DeviceDeletedTimeout: 30 * time.Millisecond})
m.Insert(&VM{
ID: "i-1",
Status: StateRunning,
Instance: &ec2.Instance{},
QMPClient: qmpClient,
EBSRequests: types.EBSRequests{
Requests: []types.EBSRequest{{Name: "vol-1", DeviceName: "/dev/sdf"}},
},
})
device, err := m.DetachVolume(t.Context(), "i-1", "vol-1", "", false)
require.NoError(t, err)
assert.Equal(t, "/dev/sdf", device)
assert.Equal(t, int32(2), blockdevAttempts.Load(),
"a missed DEVICE_DELETED event must still resolve via the bounded blockdev-del retry")
}
// TestDetachVolume_BlockdevDelAlreadyRemoved_IdempotentSuccess covers the
// AWS-CLI-retry resume path end to end: a prior detach removed the block
// node, so blockdev-del returns "Failed to find node" on the second call.
// DetachVolume must treat that as success and complete the detach, not
// surface it as a hard failure.
func TestDetachVolume_BlockdevDelAlreadyRemoved_IdempotentSuccess(t *testing.T) {
qmpClient, cancel := newMockQMPClientWithEvents(t, func(cmd qmp.QMPCommand, srv *mockQMPServer) map[string]any {
if cmd.Execute == "blockdev-del" {
return map[string]any{"error": map[string]any{"class": "GenericError", "desc": "Failed to find node with node-name='nbd-vol-1'"}}
}
return nil
})
defer cancel()
mounter := &fakeVolumeMounter{}
m := NewManagerWithDeps(Deps{VolumeMounter: mounter, DeviceDeletedTimeout: 30 * time.Millisecond})
m.Insert(&VM{
ID: "i-1",
Status: StateRunning,
Instance: &ec2.Instance{},
QMPClient: qmpClient,
EBSRequests: types.EBSRequests{
Requests: []types.EBSRequest{{Name: "vol-1", DeviceName: "/dev/sdf"}},
},
})
device, err := m.DetachVolume(t.Context(), "i-1", "vol-1", "", false)
require.NoError(t, err)
assert.Equal(t, "/dev/sdf", device)
assert.Equal(t, []string{"vol-1"}, mounter.unmountedOne,
"an already-removed block node must resume through to the unmount seal")
}