fix(ec2): drain a snapshot's volume on the node hosting it (#589)

spinifex/daemon/daemon_handlers.go

@@ -160,6 +160,13 @@ func (d *Daemon) handleEC2Events(msg *nats.Msg) {
 		d.handleAttachVolume(ctx, msg, command, instance)
 	case command.Attributes.DetachVolume:
 		d.handleDetachVolume(ctx, msg, command, instance)
+	case command.Attributes.DrainVolume:
+		// A drain flushes the volume's whole dirty set to S3, so it is bounded by
+		// the dirty set rather than by a unit of work. nats.go delivers a
+		// subscription's messages serially, so running it inline would stall
+		// every other command for this instance — stop, terminate, hot-plug —
+		// for the duration. It touches no VM state, so it is safe off-thread.
+		go d.handleDrainVolume(ctx, msg, command, instance)
 	case command.Attributes.AttachENI:
 		d.handleAttachNetworkInterface(ctx, msg, command, instance)
 	case command.Attributes.DetachENI:

spinifex/daemon/daemon_handlers_volume.go

@@ -9,6 +9,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	handlers_ec2_snapshot "github.com/mulgadc/spinifex/spinifex/handlers/ec2/snapshot"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
@@ -139,6 +140,51 @@ func (d *Daemon) handleDetachVolume(ctx context.Context, msg *nats.Msg, command
 	d.respondWithVolumeAttachment(msg, command.DetachVolumeData.VolumeID, command.ID, deviceName, "detaching")
 }
 
+// handleDrainVolume flushes the volume's in-flight writes to S3 by dialing the
+// drain socket the local NBD plugin serves. ec2.CreateSnapshot is answered by
+// any node in the cluster, but only the node hosting the instance subscribes
+// ec2.cmd.{instanceID}, so routing the drain here lands it where the writes are.
+//
+// Runs on its own goroutine (see handleEC2Events) so a long flush cannot hold
+// the instance's command subscription.
+func (d *Daemon) handleDrainVolume(ctx context.Context, msg *nats.Msg, command types.EC2InstanceCommand, instance *vm.VM) {
+	ctx, span := startOpSpan(ctx, "ec2.DrainVolume", command.ID)
+	var err error
+	defer func() { endOpSpan(span, err) }()
+
+	if command.DrainVolumeData == nil || command.DrainVolumeData.VolumeID == "" {
+		slog.ErrorContext(ctx, "DrainVolume: missing drain volume data", "instanceId", command.ID)
+		respondWithError(msg, awserrors.ErrorInvalidParameterValue)
+		return
+	}
+
+	volumeID := command.DrainVolumeData.VolumeID
+
+	// Only stable post-teardown states prove the volume was sealed. Transitional
+	// states still have a live or shutting-down writer, so they must drain or
+	// fail rather than acknowledge an older checkpoint as current.
+	status := d.vmMgr.Status(instance)
+	if status == vm.StateStopped || status == vm.StateTerminated {
+		slog.InfoContext(ctx, "DrainVolume: instance teardown is complete, nothing to drain",
+			"volumeId", volumeID, "instanceId", command.ID, "status", status)
+		respondWithJSON(msg, types.DrainVolumeResponse{VolumeID: volumeID, Status: types.DrainVolumeStatusNotRunning})
+		return
+	}
+
+	slog.InfoContext(ctx, "Draining volume before snapshot", "volumeId", volumeID, "instanceId", command.ID)
+
+	// A missing or unresponsive socket under a running instance means the writes
+	// cannot be made current, so the caller must fail its snapshot rather than
+	// read a stale checkpoint.
+	if err = handlers_ec2_snapshot.DrainVolumeSocket(d.config.DataDir, volumeID); err != nil {
+		slog.ErrorContext(ctx, "DrainVolume: drain failed", "volumeId", volumeID, "instanceId", command.ID, "err", err)
+		respondWithError(msg, awserrors.ErrorServerInternal)
+		return
+	}
+
+	respondWithJSON(msg, types.DrainVolumeResponse{VolumeID: volumeID, Status: types.DrainVolumeStatusDrained})
+}
+
 // attachDetachErrorCode maps a vm.Manager error returned by AttachVolume
 // or DetachVolume to the AWS API error code that the SDK expects.
 func attachDetachErrorCode(err error) string {

spinifex/daemon/daemon_handlers_volume_test.go

@@ -15,11 +15,15 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/qmp"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
 	"github.com/mulgadc/spinifex/spinifex/types"
+	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/spinifex/spinifex/vm"
 	"github.com/mulgadc/viperblock/viperblock"
+	"github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+	"go.uber.org/goleak"
 )
 
 // TestAttachDetachErrorCode locks down the manager-error → AWS-API-code
@@ -314,3 +318,254 @@ func TestAttachVolume_IdempotentSameInstance_DeviceMismatch(t *testing.T) {
 	require.NoError(t, err)
 	assert.Contains(t, string(resp.Data), awserrors.ErrorVolumeInUse)
 }
+
+// drainCommandFor builds the drain command a snapshot addresses to the node
+// hosting instanceID.
+func drainCommandFor(t *testing.T, instanceID, volumeID string) []byte {
+	t.Helper()
+	command := types.EC2InstanceCommand{
+		ID:              instanceID,
+		Attributes:      types.EC2CommandAttributes{DrainVolume: true},
+		DrainVolumeData: &types.DrainVolumeData{VolumeID: volumeID},
+	}
+	data, err := json.Marshal(command)
+	require.NoError(t, err)
+	return data
+}
+
+// drainTestDaemon returns a daemon owning a running instance, with a data dir
+// short enough to hold the drain socket path.
+func drainTestDaemon(t *testing.T, instanceID string) *Daemon {
+	t.Helper()
+	return drainTestDaemonInState(t, instanceID, vm.StateRunning)
+}
+
+// drainTestDaemonInState is drainTestDaemon for an instance the node still
+// holds in some other state.
+func drainTestDaemonInState(t *testing.T, instanceID string, status vm.InstanceState) *Daemon {
+	t.Helper()
+	daemon := createTestDaemon(t, sharedNATSURL)
+	daemon.config.DataDir = testutil.SocketTempDir(t)
+	daemon.vmMgr.Insert(&vm.VM{
+		ID:           instanceID,
+		Status:       status,
+		AccountID:    testAccountID,
+		InstanceType: getTestInstanceType(t),
+		Instance:     &ec2.Instance{},
+	})
+	return daemon
+}
+
+// drainRequest issues a drain on the instance's command subject, returning the
+// raw reply.
+func drainRequest(t *testing.T, daemon *Daemon, instanceID, volumeID string) *nats.Msg {
+	t.Helper()
+	reply, err := sendDrainCommand(daemon, instanceID, drainCommandFor(t, instanceID, volumeID))
+	require.NoError(t, err)
+	return reply
+}
+
+// sendDrainCommand is drainRequest without the assertions, so a test can issue
+// one from a goroutine and assert on the result back on the test goroutine.
+func sendDrainCommand(daemon *Daemon, instanceID string, command []byte) (*nats.Msg, error) {
+	msg := nats.NewMsg("ec2.cmd." + instanceID)
+	msg.Data = command
+	msg.Header.Set(utils.AccountIDHeader, testAccountID)
+	return daemon.natsConn.RequestMsg(msg, 5*time.Second)
+}
+
+// subscribeEC2Commands wires the daemon's real dispatcher onto the per-instance
+// command subject, so tests exercise the same serial delivery production has.
+func subscribeEC2Commands(t *testing.T, daemon *Daemon, instanceID string) {
+	t.Helper()
+	sub, err := daemon.natsConn.Subscribe("ec2.cmd."+instanceID, daemon.handleEC2Events)
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = sub.Unsubscribe() })
+}
+
+// A drain routed to the node hosting the instance reaches the local socket and
+// acks, which is what lets a snapshot answered elsewhere read a current
+// checkpoint.
+func TestDrainVolume_HostNodeAcksLocalSocket(t *testing.T) {
+	const instanceID, volumeID = "i-drain-ok", "vol-drain-ok"
+	daemon := drainTestDaemon(t, instanceID)
+	testutil.StartDrainSocket(t, daemon.config.DataDir, volumeID, "OK\n")
+
+	resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+		testAccountID, drainCommandFor(t, instanceID, volumeID))
+
+	var ack types.DrainVolumeResponse
+	require.NoError(t, json.Unmarshal(resp.Data, &ack))
+	assert.Equal(t, volumeID, ack.VolumeID)
+	assert.Equal(t, types.DrainVolumeStatusDrained, ack.Status)
+}
+
+// The owning node having no socket for the volume means the writes cannot be
+// flushed. That must reach the caller as an error, never as a silent success:
+// the caller would otherwise snapshot a stale checkpoint.
+func TestDrainVolume_HostNodeWithoutSocketFails(t *testing.T) {
+	const instanceID, volumeID = "i-drain-nosock", "vol-drain-nosock"
+	daemon := drainTestDaemon(t, instanceID)
+
+	resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+		testAccountID, drainCommandFor(t, instanceID, volumeID))
+
+	assert.Equal(t, awserrors.ErrorServerInternal, replyErrCode(t, resp.Data))
+}
+
+// A plugin that refuses the drain (ERR rather than OK) is a failure, not an ack.
+func TestDrainVolume_HostNodeRelaysSocketFailure(t *testing.T) {
+	const instanceID, volumeID = "i-drain-err", "vol-drain-err"
+	daemon := drainTestDaemon(t, instanceID)
+	testutil.StartDrainSocket(t, daemon.config.DataDir, volumeID, "ERR\n")
+
+	resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+		testAccountID, drainCommandFor(t, instanceID, volumeID))
+
+	assert.Equal(t, awserrors.ErrorServerInternal, replyErrCode(t, resp.Data))
+}
+
+// A drain command with no volume is a caller error, not a server fault.
+func TestDrainVolume_MissingVolumeDataIsInvalidParameter(t *testing.T) {
+	const instanceID = "i-drain-novol"
+	daemon := drainTestDaemon(t, instanceID)
+
+	command := types.EC2InstanceCommand{
+		ID:         instanceID,
+		Attributes: types.EC2CommandAttributes{DrainVolume: true},
+	}
+	data, err := json.Marshal(command)
+	require.NoError(t, err)
+
+	resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+		testAccountID, data)
+
+	assert.Equal(t, awserrors.ErrorInvalidParameterValue, replyErrCode(t, resp.Data))
+}
+
+// Only the instance's owner may drain its volume: the command carries the same
+// ownership gate as every other per-instance command.
+func TestDrainVolume_ForeignAccountRejected(t *testing.T) {
+	const instanceID, volumeID = "i-drain-foreign", "vol-drain-foreign"
+	daemon := drainTestDaemon(t, instanceID)
+	testutil.StartDrainSocket(t, daemon.config.DataDir, volumeID, "OK\n")
+
+	resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+		"999988887777", drainCommandFor(t, instanceID, volumeID))
+
+	assert.Equal(t, awserrors.ErrorInvalidInstanceIDNotFound, replyErrCode(t, resp.Data))
+}
+
+// Only stable post-teardown states prove the volume is sealed. These states
+// acknowledge not-running without touching the absent plugin socket.
+func TestDrainVolume_CompletedTeardownAcksNotRunning(t *testing.T) {
+	for _, status := range []vm.InstanceState{vm.StateStopped, vm.StateTerminated} {
+		t.Run(string(status), func(t *testing.T) {
+			instanceID := "i-drain-" + string(status)
+			volumeID := "vol-drain-" + string(status)
+			daemon := drainTestDaemonInState(t, instanceID, status)
+
+			resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+				testAccountID, drainCommandFor(t, instanceID, volumeID))
+
+			var ack types.DrainVolumeResponse
+			require.NoError(t, json.Unmarshal(resp.Data, &ack))
+			assert.Equal(t, volumeID, ack.VolumeID)
+			assert.Equal(t, types.DrainVolumeStatusNotRunning, ack.Status)
+		})
+	}
+}
+
+// Transitional and uncertain states are not proof of a sealed checkpoint. An
+// absent socket in any of them must fail rather than silently snapshot stale
+// data while launch or teardown may still own the volume.
+func TestDrainVolume_UnsealedStatesRequireDrain(t *testing.T) {
+	statuses := []vm.InstanceState{
+		vm.StateRunning,
+		vm.StateStopping,
+		vm.StateShuttingDown,
+		vm.StatePending,
+		vm.StateProvisioning,
+		vm.StateError,
+	}
+	for _, status := range statuses {
+		t.Run(string(status), func(t *testing.T) {
+			instanceID := "i-drain-" + string(status)
+			volumeID := "vol-drain-" + string(status)
+			daemon := drainTestDaemonInState(t, instanceID, status)
+
+			resp := requestHandler(t, daemon.natsConn, "ec2.cmd."+instanceID, daemon.handleEC2Events,
+				testAccountID, drainCommandFor(t, instanceID, volumeID))
+
+			assert.Equal(t, awserrors.ErrorServerInternal, replyErrCode(t, resp.Data))
+		})
+	}
+}
+
+// nats.go delivers a subscription's messages serially, so a drain run inline
+// would hold ec2.cmd.{instanceID} for the length of the flush — stalling stop,
+// terminate and hot-plug for that instance. A drain still flushing must not
+// keep the next command waiting.
+func TestDrainVolume_SlowDrainDoesNotBlockTheCommandSubject(t *testing.T) {
+	const instanceID = "i-drain-parallel"
+	daemon := drainTestDaemon(t, instanceID)
+	accepted, release := testutil.StartBlockingDrainSocket(t, daemon.config.DataDir, "vol-drain-slow", "OK\n")
+	testutil.StartDrainSocket(t, daemon.config.DataDir, "vol-drain-fast", "OK\n")
+	subscribeEC2Commands(t, daemon, instanceID)
+
+	type reply struct {
+		msg *nats.Msg
+		err error
+	}
+	slowCommand := drainCommandFor(t, instanceID, "vol-drain-slow")
+	slow := make(chan reply, 1)
+	go func() {
+		msg, err := sendDrainCommand(daemon, instanceID, slowCommand)
+		slow <- reply{msg, err}
+	}()
+
+	// Only assert once the slow drain is genuinely mid-flush.
+	select {
+	case <-accepted:
+	case <-time.After(5 * time.Second):
+		t.Fatal("the slow drain never reached its socket")
+	}
+
+	fast := drainRequest(t, daemon, instanceID, "vol-drain-fast")
+	var ack types.DrainVolumeResponse
+	require.NoError(t, json.Unmarshal(fast.Data, &ack))
+	assert.Equal(t, types.DrainVolumeStatusDrained, ack.Status)
+
+	release()
+	select {
+	case got := <-slow:
+		require.NoError(t, got.err)
+		require.NoError(t, json.Unmarshal(got.msg.Data, &ack))
+		assert.Equal(t, types.DrainVolumeStatusDrained, ack.Status)
+	case <-time.After(5 * time.Second):
+		t.Fatal("the released drain never replied")
+	}
+}
+
+// Running the drain off the delivery goroutine must not leak it: the goroutine
+// has to be gone once the ack is on the wire, or a snapshot loop against a
+// wedged plugin accumulates them.
+func TestDrainVolume_DispatchedGoroutineDoesNotLeak(t *testing.T) {
+	const instanceID, volumeID = "i-drain-leak", "vol-drain-leak"
+	daemon := drainTestDaemon(t, instanceID)
+	testutil.StartDrainSocket(t, daemon.config.DataDir, volumeID, "OK\n")
+	subscribeEC2Commands(t, daemon, instanceID)
+
+	// Take the baseline after one full round trip: the first request is what
+	// makes nats.go stand up its long-lived response-inbox goroutine, which
+	// would otherwise read as the leak.
+	var ack types.DrainVolumeResponse
+	require.NoError(t, json.Unmarshal(drainRequest(t, daemon, instanceID, volumeID).Data, &ack))
+	require.Equal(t, types.DrainVolumeStatusDrained, ack.Status)
+	ignoreExisting := goleak.IgnoreCurrent()
+
+	require.NoError(t, json.Unmarshal(drainRequest(t, daemon, instanceID, volumeID).Data, &ack))
+	require.Equal(t, types.DrainVolumeStatusDrained, ack.Status)
+
+	goleak.VerifyNone(t, ignoreExisting)
+}

spinifex/handlers/NATS_HANDLER_NAMING.md

@@ -54,6 +54,27 @@ handleEC2CreateSnapshot     → ec2.CreateSnapshot       // Create volume snapsh
 handleEC2DeleteSnapshot     → ec2.DeleteSnapshot       // Remove snapshot
 ```
 
+### Per-instance commands (`ec2.cmd.{instanceID}`)
+
+Actions that must run on the node hosting a VM do not use the `spinifex-workers` queue group, which answers on any node. They travel on `ec2.cmd.{instanceID}`, subscribed only by the node running that instance, so the command is self-routing: the publisher never resolves instance → node.
+
+The subject carries a single `types.EC2InstanceCommand`, dispatched by `handleEC2Events` on the boolean set in `EC2CommandAttributes`:
+
+```go
+AttachVolume, DetachVolume    // QMP hot-plug of an EBS volume
+DrainVolume                   // flush the volume's in-flight writes to S3
+AttachENI, DetachENI          // ENI hot-plug (awsvpc tasks)
+StartInstance, StopInstance, TerminateInstance, RebootInstance
+AssociateIamInstanceProfile, SetSpotLineage
+SetInstanceTags, RemoveInstanceTags
+```
+
+`DrainVolume` exists because `ec2.CreateSnapshot` is queue-grouped but the drain socket that flushes in-flight writes is served by the NBD plugin on the node hosting the volume. The node answering the snapshot routes the drain here and waits for the ack before reading the live checkpoint; a volume with a live writer that does not ack fails the snapshot rather than capturing stale data.
+
+Establishing that there *is* a live writer is the subject's job, because the attachment record alone does not say so — stop deliberately leaves a boot volume attached while tearing down both the plugin and this subscription. Two replies therefore mean "nothing is writing, the sealed checkpoint stands" rather than failure: `DrainVolumeStatusNotRunning`, from a node that still holds the instance but is not running it, and `nats.ErrNoResponders`, meaning no node runs it at all.
+
+The drain is dispatched on its own goroutine. It is bounded by the volume's dirty set, and nats.go delivers a subscription's messages serially, so running it inline would stall every other command for that instance for the length of the flush.
+
 ### VPC Networking
 ```go
 handleEC2CreateVpc          → ec2.CreateVpc            // Create VPC

spinifex/handlers/ec2/snapshot/drain_test.go

@@ -0,0 +1,316 @@
+package handlers_ec2_snapshot
+
+import (
+	"context"
+	"encoding/json"
+	"net/http"
+	"net/http/httptest"
+	"testing"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/ec2"
+	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/config"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate"
+	"github.com/mulgadc/spinifex/spinifex/objectstore"
+	"github.com/mulgadc/spinifex/spinifex/testutil"
+	"github.com/mulgadc/spinifex/spinifex/types"
+	"github.com/mulgadc/spinifex/spinifex/utils"
+	"github.com/mulgadc/viperblock/viperblock"
+	"github.com/nats-io/nats.go"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const drainInstanceID = "i-drain-host"
+
+// setupDrainService returns a snapshot service configured like a daemon that
+// answers ec2.CreateSnapshot: a reachable Predastore host (so drains are not
+// skipped as metadata-only), a data dir with no drain socket in it, and a NATS
+// connection. The Predastore stub rejects everything, so a snapshot that gets
+// past the drain fails fast rather than hanging on a dead address.
+func setupDrainService(t *testing.T) (*SnapshotServiceImpl, *objectstore.MemoryObjectStore, *nats.Conn) {
+	t.Helper()
+
+	srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
+		w.WriteHeader(http.StatusForbidden)
+	}))
+	t.Cleanup(srv.Close)
+
+	_, nc := testutil.StartTestNATS(t)
+	store := objectstore.NewMemoryObjectStore()
+	cfg := &config.Config{
+		DataDir: testutil.SocketTempDir(t),
+		Predastore: config.PredastoreConfig{
+			Host:   srv.URL,
+			Bucket: "test-bucket",
+		},
+	}
+	return NewSnapshotServiceImplWithStore(cfg, store, nc), store, nc
+}
+
+// seedVolumeAttachment writes the control-plane state.json that decides whether
+// a snapshot must drain.
+func seedVolumeAttachment(t *testing.T, store *objectstore.MemoryObjectStore, volumeID, state, instanceID string) {
+	t.Helper()
+	require.NoError(t, volumestate.Write(context.Background(), store, "test-bucket", volumeID, volumestate.Record{
+		State:            state,
+		AttachedInstance: instanceID,
+	}))
+}
+
+// drainResponder answers the per-instance command subject the way the node
+// hosting the instance does, and hands each request it received to the test.
+func drainResponder(t *testing.T, nc *nats.Conn, instanceID string, reply []byte) chan *nats.Msg {
+	t.Helper()
+
+	got := make(chan *nats.Msg, 4)
+	sub, err := nc.Subscribe("ec2.cmd."+instanceID, func(msg *nats.Msg) {
+		got <- msg
+		_ = msg.Respond(reply)
+	})
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = sub.Unsubscribe() })
+
+	return got
+}
+
+// drainedAck is the reply of a node whose drain reached S3.
+func drainedAck(t *testing.T, volumeID string) []byte {
+	t.Helper()
+	data, err := json.Marshal(types.DrainVolumeResponse{VolumeID: volumeID, Status: types.DrainVolumeStatusDrained})
+	require.NoError(t, err)
+	return data
+}
+
+// requireReturnsWithin fails the test if fn has not returned successfully by
+// limit, so a drain that blocks on a socket it should not have dialed reports
+// as a failure rather than as a slow test.
+func requireReturnsWithin(t *testing.T, limit time.Duration, fn func() error) {
+	t.Helper()
+
+	done := make(chan error, 1)
+	go func() { done <- fn() }()
+
+	select {
+	case err := <-done:
+		require.NoError(t, err)
+	case <-time.After(limit):
+		t.Fatalf("drain did not return within %s", limit)
+	}
+}
+
+// awaitDrainCommand returns the command the hosting node received, failing the
+// test if none arrives.
+func awaitDrainCommand(t *testing.T, got chan *nats.Msg) types.EC2InstanceCommand {
+	t.Helper()
+	select {
+	case msg := <-got:
+		var command types.EC2InstanceCommand
+		require.NoError(t, json.Unmarshal(msg.Data, &command))
+		assert.Equal(t, testAccountID, msg.Header.Get(utils.AccountIDHeader))
+		return command
+	case <-time.After(2 * time.Second):
+		t.Fatal("the node hosting the volume never received a drain command")
+		return types.EC2InstanceCommand{}
+	}
+}
+
+// An attached volume is drained on the node hosting its instance, which is the
+// only node with the socket — the node answering CreateSnapshot usually is not.
+func TestDrainVolume_AttachedRoutesToHostingNode(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-attached", "in-use", drainInstanceID)
+	got := drainResponder(t, nc, drainInstanceID, drainedAck(t, "vol-attached"))
+
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-attached", viperblock.VolumeMetadata{}, testAccountID))
+
+	command := awaitDrainCommand(t, got)
+	assert.True(t, command.Attributes.DrainVolume)
+	assert.Equal(t, drainInstanceID, command.ID)
+	require.NotNil(t, command.DrainVolumeData)
+	assert.Equal(t, "vol-attached", command.DrainVolumeData.VolumeID)
+}
+
+// No responder on the command subject means no node is running the instance,
+// which is what a stopped instance looks like: stop tears down both the plugin
+// and the subscription but deliberately leaves a boot volume attached. Failing
+// here would make every stopped instance's root volume unsnapshottable.
+func TestDrainVolume_AttachedToStoppedInstanceTakesStoppedPath(t *testing.T) {
+	svc, store, _ := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-no-host", "in-use", drainInstanceID)
+
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-no-host", viperblock.VolumeMetadata{}, testAccountID))
+}
+
+// A host that still holds the instance but reports it not running has nothing
+// to flush, so the sealed checkpoint stands. This is the host-drain stop, which
+// unmounts the volumes but keeps the subscription.
+func TestDrainVolume_NotRunningAckTakesStoppedPath(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-not-running", "in-use", drainInstanceID)
+	ack, err := json.Marshal(types.DrainVolumeResponse{
+		VolumeID: "vol-not-running", Status: types.DrainVolumeStatusNotRunning,
+	})
+	require.NoError(t, err)
+	got := drainResponder(t, nc, drainInstanceID, ack)
+
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-not-running", viperblock.VolumeMetadata{}, testAccountID))
+	assert.True(t, awaitDrainCommand(t, got).Attributes.DrainVolume)
+}
+
+// A hosting node that reports the drain failed (wedged plugin, socket gone
+// mid-restart) fails the snapshot rather than falling through.
+func TestDrainVolume_AttachedWithFailedAckFails(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-drain-err", "in-use", drainInstanceID)
+	drainResponder(t, nc, drainInstanceID,
+		[]byte(`{"Code":"`+awserrors.ErrorServerInternal+`","Message":"drain failed"}`))
+
+	err := svc.drainVolume(context.Background(), "vol-drain-err", viperblock.VolumeMetadata{}, testAccountID)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "drain failed")
+}
+
+// An ack that is not the drained ack is not evidence the writes reached S3.
+func TestDrainVolume_AttachedWithUnexpectedAckFails(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-odd-ack", "in-use", drainInstanceID)
+	drainResponder(t, nc, drainInstanceID, []byte(`{}`))
+
+	err := svc.drainVolume(context.Background(), "vol-odd-ack", viperblock.VolumeMetadata{}, testAccountID)
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "unexpected ack")
+}
+
+// Nothing writes to an unattached volume, so the checkpoint its Close() left
+// behind is current and no drain is attempted — not even locally. Dialing the
+// socket is not a probe: it makes the plugin run a full flush, so an available
+// volume must be decided from the record without touching it.
+func TestDrainVolume_AvailableDoesNotDialLocalSocket(t *testing.T) {
+	svc, store, _ := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-available", "available", "")
+
+	// A socket that never answers: a dial would park here until the read
+	// deadline, so returning promptly is what proves none was made.
+	testutil.StartBlockingDrainSocket(t, svc.config.DataDir, "vol-available", "OK\n")
+
+	// A nil connection proves no command was issued: routing one would error.
+	svc.natsConn = nil
+
+	requireReturnsWithin(t, 2*time.Second, func() error {
+		return svc.drainVolume(context.Background(), "vol-available", viperblock.VolumeMetadata{}, testAccountID)
+	})
+}
+
+// A volume recorded as in-use but with no instance (drifted state) has no node
+// to address, so it takes the stopped path rather than failing every snapshot.
+func TestDrainVolume_InUseWithoutInstanceTakesStoppedPath(t *testing.T) {
+	svc, store, _ := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-no-instance", "in-use", "")
+	svc.natsConn = nil
+
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-no-instance", viperblock.VolumeMetadata{}, testAccountID))
+}
+
+// When the volume is served by this node the socket is local, so the drain
+// short-circuits without a hop. This is every single-node deployment.
+func TestDrainVolume_LocalSocketShortCircuits(t *testing.T) {
+	svc, store, _ := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-local", "in-use", drainInstanceID)
+	testutil.StartDrainSocket(t, svc.config.DataDir, "vol-local", "OK\n")
+
+	// A nil connection proves the local ack was enough: no responder exists for
+	// drainInstanceID, so a routed command would fail.
+	svc.natsConn = nil
+
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-local", viperblock.VolumeMetadata{}, testAccountID))
+}
+
+// A socket that answers without acking means this node does serve the volume
+// and the flush failed. Routing that onward would come back to this same node
+// and re-run the drain that has already failed, so it fails here instead.
+func TestDrainVolume_LocalSocketErrorAckFailsWithoutRouting(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-local-err", "in-use", drainInstanceID)
+	testutil.StartDrainSocket(t, svc.config.DataDir, "vol-local-err", "ERR\n")
+	got := drainResponder(t, nc, drainInstanceID, drainedAck(t, "vol-local-err"))
+
+	err := svc.drainVolume(context.Background(), "vol-local-err", viperblock.VolumeMetadata{}, testAccountID)
+	require.Error(t, err)
+
+	assert.Empty(t, got, "a local drain that failed must not be re-issued to the host node")
+}
+
+// state.json is the control-plane-owned attachment record; the copy inside
+// config.json is rewritten by the live NBD plugin from its stale in-memory
+// state, so a volume reading "available" there can still be attached.
+func TestDrainVolume_StateJSONOverridesConfigJSON(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-stale-config", "in-use", drainInstanceID)
+	got := drainResponder(t, nc, drainInstanceID, drainedAck(t, "vol-stale-config"))
+
+	stale := viperblock.VolumeMetadata{State: "available"}
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-stale-config", stale, testAccountID))
+
+	assert.True(t, awaitDrainCommand(t, got).Attributes.DrainVolume)
+}
+
+// Volumes predating the state.json split have no state object, so the state
+// embedded in config.json is all there is to decide from.
+func TestDrainVolume_FallsBackToConfigJSONWhenNoStateObject(t *testing.T) {
+	svc, _, nc := setupDrainService(t)
+	got := drainResponder(t, nc, drainInstanceID, drainedAck(t, "vol-legacy"))
+
+	legacy := viperblock.VolumeMetadata{State: "in-use", AttachedInstance: drainInstanceID}
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-legacy", legacy, testAccountID))
+
+	assert.True(t, awaitDrainCommand(t, got).Attributes.DrainVolume)
+}
+
+// Without Predastore the snapshot is metadata-only and never reads a live
+// checkpoint, so there is nothing for a drain to make current.
+func TestDrainVolume_MetadataOnlySkipsDrain(t *testing.T) {
+	svc, store, _ := setupDrainService(t)
+	seedVolumeAttachment(t, store, "vol-meta-only", "in-use", drainInstanceID)
+	svc.config.Predastore.Host = ""
+	svc.natsConn = nil
+
+	require.NoError(t, svc.drainVolume(context.Background(), "vol-meta-only", viperblock.VolumeMetadata{}, testAccountID))
+}
+
+// CreateSnapshot drains before it reads the checkpoint. The snapshot itself
+// then fails against the rejecting Predastore stub; what is asserted here is
+// that the drain was routed to the volume's host first.
+func TestCreateSnapshot_DrainsAttachedVolume(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	createTestVolume(t, store, "vol-snap-drain", 10)
+	seedVolumeAttachment(t, store, "vol-snap-drain", "in-use", drainInstanceID)
+	got := drainResponder(t, nc, drainInstanceID, drainedAck(t, "vol-snap-drain"))
+
+	_, err := svc.CreateSnapshot(context.Background(),
+		&ec2.CreateSnapshotInput{VolumeId: aws.String("vol-snap-drain")}, testAccountID)
+	require.Error(t, err)
+
+	command := awaitDrainCommand(t, got)
+	assert.True(t, command.Attributes.DrainVolume)
+	assert.Equal(t, "vol-snap-drain", command.DrainVolumeData.VolumeID)
+}
+
+// An attached volume whose host reports the drain failed fails the snapshot: a
+// retryable error is strictly better than a successful snapshot of stale data.
+func TestCreateSnapshot_UndrainableAttachedVolumeFails(t *testing.T) {
+	svc, store, nc := setupDrainService(t)
+	createTestVolume(t, store, "vol-snap-nodrain", 10)
+	seedVolumeAttachment(t, store, "vol-snap-nodrain", "in-use", drainInstanceID)
+	got := drainResponder(t, nc, drainInstanceID,
+		[]byte(`{"Code":"`+awserrors.ErrorServerInternal+`","Message":"drain failed"}`))
+
+	_, err := svc.CreateSnapshot(context.Background(),
+		&ec2.CreateSnapshotInput{VolumeId: aws.String("vol-snap-nodrain")}, testAccountID)
+	require.Error(t, err)
+	assert.Equal(t, awserrors.ErrorServerInternal, err.Error())
+
+	assert.True(t, awaitDrainCommand(t, got).Attributes.DrainVolume)
+}

spinifex/handlers/ec2/snapshot/service_impl.go

@@ -22,9 +22,11 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate"
 	"github.com/mulgadc/spinifex/spinifex/kvutil"
 	"github.com/mulgadc/spinifex/spinifex/migrate"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
+	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/utils"
 	"github.com/mulgadc/viperblock/viperblock"
 	vbs3 "github.com/mulgadc/viperblock/viperblock/backends/s3"
@@ -243,10 +245,19 @@ func (s *SnapshotServiceImpl) CreateSnapshot(ctx context.Context, input *ec2.Cre
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
+	// Flush the writes still buffered by whichever node serves the volume, so the
+	// live checkpoint this snapshot is about to read is current. An attached
+	// volume that cannot be drained fails here rather than silently producing a
+	// snapshot of an older checkpoint.
+	if err := s.drainVolume(ctx, volumeID, volumeConfig.VolumeMetadata, accountID); err != nil {
+		slog.ErrorContext(ctx, "CreateSnapshot: drain failed", "volumeId", volumeID, "snapshotId", snapshotID, "err", err)
+		return nil, errors.New(awserrors.ErrorServerInternal)
+	}
+
 	// Snapshot the viperblock volume by reading the live checkpoint from S3.
 	// The live checkpoint is updated on every NBD Flush by the running nbdkit process.
 	// If the volume is not mounted (stopped), LoadLiveCheckpoint falls back to the
-	// numbered checkpoint written by Close. No IPC with nbdkit required.
+	// numbered checkpoint written by Close.
 	if err := s.snapshotVolume(volumeID, snapshotID, volumeConfig.VolumeMetadata.SizeGiB*1024*1024*1024); err != nil {
 		slog.ErrorContext(ctx, "CreateSnapshot: viperblock snapshot failed", "volumeId", volumeID, "snapshotId", snapshotID, "err", err)
 		return nil, errors.New(awserrors.ErrorServerInternal)
@@ -292,9 +303,172 @@ func (s *SnapshotServiceImpl) CreateSnapshot(ctx context.Context, input *ec2.Cre
 	return snapshotConfigToEC2(snapshotCfg), nil
 }
 
+const (
+	// drainDialTimeout bounds the connect to the local drain socket. The socket
+	// is either being served on this node or it is not, so a slow connect means
+	// a wedged plugin rather than a busy one.
+	drainDialTimeout = time.Second
+
+	// drainAckTimeout bounds DrainToBackend on the hosting node: it flushes the
+	// WAL and the live checkpoint to S3, so it scales with the dirty set rather
+	// than with a fixed unit of work.
+	drainAckTimeout = 30 * time.Second
+
+	// drainRequestTimeout bounds the ec2.cmd round-trip. It exceeds
+	// drainAckTimeout so a slow drain surfaces as the hosting node's error
+	// instead of an opaque NATS timeout here.
+	drainRequestTimeout = 35 * time.Second
+)
+
+// ErrNoDrainSocket reports that this node has no drain socket for the volume,
+// which is how a node learns it does not serve it. It is deliberately distinct
+// from a socket that answered badly: only the former is worth routing onward,
+// and conflating the two makes a node re-run a drain that has already failed.
+var ErrNoDrainSocket = errors.New("no drain socket on this node")
+
+// DrainVolumeSocket asks the NBD plugin serving volumeID on this node to flush
+// its WAL chunks and live checkpoint to S3, then waits for the ack. The plugin
+// creates the socket under {dataDir}/viperblock/{volumeID}/, so this only
+// succeeds on the node hosting the volume; the daemon calls it when a drain
+// command is routed to it.
+func DrainVolumeSocket(dataDir, volumeID string) error {
+	sockPath := filepath.Join(dataDir, "viperblock", volumeID, "snapshot.sock")
+	conn, err := net.DialTimeout("unix", sockPath, drainDialTimeout)
+	if err != nil {
+		return fmt.Errorf("%w: dial %s: %w", ErrNoDrainSocket, sockPath, err)
+	}
+	defer conn.Close()
+
+	_ = conn.SetDeadline(time.Now().Add(drainAckTimeout))
+	buf := make([]byte, 8)
+	n, err := conn.Read(buf)
+	if err != nil {
+		return fmt.Errorf("read drain ack for %s: %w", volumeID, err)
+	}
+	if !strings.HasPrefix(string(buf[:n]), "OK") {
+		return fmt.Errorf("drain of %s did not ack OK: %q", volumeID, string(buf[:n]))
+	}
+	return nil
+}
+
+// drainVolume flushes the volume's in-flight writes to S3 before the snapshot
+// reads the live checkpoint from there.
+//
+// The drain socket lives on the node serving the volume, but ec2.CreateSnapshot
+// is queue-grouped across every node, so the node answering the request usually
+// is not that node. Whether a drain is required is therefore decided from the
+// volume's attachment record, never from whether the socket happens to be
+// local: an attached volume is being written to right now, and snapshotting it
+// without draining silently captures an older checkpoint.
+func (s *SnapshotServiceImpl) drainVolume(ctx context.Context, volumeID string, meta viperblock.VolumeMetadata, accountID string) error {
+	// A metadata-only snapshot never reads the live checkpoint, so there is
+	// nothing for a drain to make current.
+	if s.config == nil || s.config.Predastore.Host == "" {
+		return nil
+	}
+
+	state, instanceID, err := s.volumeAttachment(ctx, volumeID, meta)
+	if err != nil {
+		return err
+	}
+
+	// Nothing is writing to an unattached volume: the checkpoint its Close()
+	// left behind is the current one. Decide this before touching the socket —
+	// dialing it is not a probe, it makes the plugin run a full flush.
+	if state != "in-use" || instanceID == "" {
+		slog.InfoContext(ctx, "drainVolume: volume not attached, snapshotting the checkpoint left by Close (stopped instance path)",
+			"volumeId", volumeID, "state", state, "attachedInstance", instanceID)
+		return nil
+	}
+
+	// Fast path: the volume is served by this node, so no hop is needed. This is
+	// every single-node deployment, and the node hosting the instance on a
+	// cluster.
+	localErr := DrainVolumeSocket(s.config.DataDir, volumeID)
+	if localErr == nil {
+		return nil
+	}
+
+	// A socket that answered but did not ack means this node does serve the
+	// volume and the flush itself failed. Routing that onward would land back on
+	// this same node and repeat the drain that has already failed.
+	if !errors.Is(localErr, ErrNoDrainSocket) {
+		return fmt.Errorf("drain volume %s on this node: %w", volumeID, localErr)
+	}
+
+	slog.InfoContext(ctx, "drainVolume: volume attached, routing drain to the node hosting it",
+		"volumeId", volumeID, "instanceId", instanceID)
+	return s.drainOnHostNode(ctx, volumeID, instanceID, accountID)
+}
+
+// volumeAttachment returns the volume's authoritative state and attached
+// instance. state.json is control-plane-owned and authoritative; the copy in
+// config.json (passed in as meta) is rewritten by the live NBD plugin from its
+// stale in-memory state and is only used for volumes predating the split.
+func (s *SnapshotServiceImpl) volumeAttachment(ctx context.Context, volumeID string, meta viperblock.VolumeMetadata) (state, instanceID string, err error) {
+	rec, found, err := volumestate.Read(ctx, s.store, s.config.Predastore.Bucket, volumeID)
+	if err != nil {
+		return "", "", fmt.Errorf("read volume state for %s: %w", volumeID, err)
+	}
+	if found {
+		return rec.State, rec.AttachedInstance, nil
+	}
+	return meta.State, meta.AttachedInstance, nil
+}
+
+// drainOnHostNode issues the drain on the node hosting instanceID. Only that
+// node subscribes ec2.cmd.{instanceID}, so the command is self-routing and no
+// volume-to-node resolution is needed here.
+//
+// An attachment record is not proof of a live writer: stop deliberately leaves
+// a boot volume attached (daemon/vm_adapters.go's volumeMounterAdapter.Unmount)
+// while tearing down both the NBD plugin and this subscription. The two
+// not-running signals below are therefore the stopped-instance path, not a
+// failure — treating them as one would make every stopped instance's root
+// volume permanently unsnapshottable.
+func (s *SnapshotServiceImpl) drainOnHostNode(ctx context.Context, volumeID, instanceID, accountID string) error {
+	command := types.EC2InstanceCommand{
+		ID:              instanceID,
+		Attributes:      types.EC2CommandAttributes{DrainVolume: true},
+		DrainVolumeData: &types.DrainVolumeData{VolumeID: volumeID},
+	}
+
+	resp, err := utils.NATSRequest[types.DrainVolumeResponse](ctx, s.natsConn,
+		"ec2.cmd."+instanceID, command, drainRequestTimeout, accountID)
+	if err != nil {
+		// No subscriber at all: the instance runs nowhere in the cluster, so it
+		// was stopped (stop migrates it to shared KV before unsubscribing) and
+		// its volumes were sealed by Close. Warn rather than Info because the
+		// same shape appears if the hosting node has lost NATS while its VM
+		// keeps writing, which yields a stale snapshot as it did before routing
+		// existed. A live host that answers and fails is still a hard error.
+		if errors.Is(err, nats.ErrNoResponders) {
+			slog.WarnContext(ctx, "drainVolume: no node hosts the instance, treating the volume as stopped and snapshotting its sealed checkpoint",
+				"volumeId", volumeID, "instanceId", instanceID)
+			return nil
+		}
+		return fmt.Errorf("drain volume %s on the node hosting %s: %w", volumeID, instanceID, err)
+	}
+
+	switch resp.Status {
+	case types.DrainVolumeStatusDrained:
+		return nil
+	// The host still holds the instance but it is not running: nothing is
+	// writing, so there is nothing to flush. This is a host-drain stop, which
+	// keeps the VM (and this subscription) in place after unmounting.
+	case types.DrainVolumeStatusNotRunning:
+		slog.InfoContext(ctx, "drainVolume: instance is not running on its host, snapshotting the volume's sealed checkpoint",
+			"volumeId", volumeID, "instanceId", instanceID)
+		return nil
+	default:
+		return fmt.Errorf("drain volume %s on the node hosting %s: unexpected ack %q", volumeID, instanceID, resp.Status)
+	}
+}
+
 // snapshotVolume opens a read-only viperblock instance, reads the live checkpoint from S3
-// (written by nbdkit on every NBD Flush), and calls CreateSnapshot. Falls back to the
-// numbered checkpoint from Close if no live checkpoint exists (stopped volume path).
+// (written by nbdkit on every NBD Flush, and by the drain the caller has already run),
+// and calls CreateSnapshot. Falls back to the numbered checkpoint from Close if no live
+// checkpoint exists (stopped volume path).
 // If Predastore is not configured the snapshot proceeds as metadata-only.
 func (s *SnapshotServiceImpl) snapshotVolume(volumeID, snapshotID string, volumeSize uint64) error {
 	if s.config == nil || s.config.Predastore.Host == "" {
@@ -342,25 +516,6 @@ func (s *SnapshotServiceImpl) snapshotVolume(volumeID, snapshotID string, volume
 		return fmt.Errorf("load state: %w", err)
 	}
 
-	// Signal a running viperblock instance to flush WAL chunks and the live
-	// checkpoint to S3 before we read them. If the socket is absent the instance
-	// is stopped and its Close() already drained.
-	// DataDir is always set (from data_dir in spinifex.toml); viperblock volumes
-	// live under {DataDir}/viperblock/{volumeID}/ which is where the NBD plugin
-	// creates the socket.
-	sockPath := filepath.Join(s.config.DataDir, "viperblock", volumeID, "snapshot.sock")
-	if conn, err := net.DialTimeout("unix", sockPath, time.Second); err == nil {
-		_ = conn.SetDeadline(time.Now().Add(30 * time.Second))
-		buf := make([]byte, 8)
-		n, _ := conn.Read(buf)
-		conn.Close()
-		if !strings.HasPrefix(string(buf[:n]), "OK") {
-			slog.Warn("snapshotVolume: drain did not ack OK", "volumeId", volumeID, "resp", string(buf[:n]))
-		}
-	} else {
-		slog.Debug("snapshotVolume: no snapshot socket, proceeding (stopped instance path)", "volumeId", volumeID)
-	}
-
 	if err := vb.LoadLiveCheckpoint(); err != nil {
 		return fmt.Errorf("load live checkpoint: %w", err)
 	}

spinifex/handlers/ec2/volume/config_route_invariants_test.go

@@ -8,6 +8,7 @@ import (
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/viperblock/viperblock"
 	"github.com/stretchr/testify/assert"
@@ -119,7 +120,7 @@ func TestGetVolumeConfig_FallsBackWhenNoStateJSON(t *testing.T) {
 }
 
 // getStoredState reads and decodes the raw state.json from the memory store.
-func getStoredState(t *testing.T, store *objectstore.MemoryObjectStore, volumeID string) volumeStateRecord {
+func getStoredState(t *testing.T, store *objectstore.MemoryObjectStore, volumeID string) volumestate.Record {
 	t.Helper()
 	res, err := store.GetObject(context.Background(), &s3.GetObjectInput{
 		Bucket: aws.String("test-bucket"),
@@ -127,7 +128,7 @@ func getStoredState(t *testing.T, store *objectstore.MemoryObjectStore, volumeID
 	})
 	require.NoError(t, err)
 	defer res.Body.Close()
-	var rec volumeStateRecord
+	var rec volumestate.Record
 	require.NoError(t, json.NewDecoder(res.Body).Decode(&rec))
 	return rec
 }

spinifex/handlers/ec2/volume/lifecycle_invariants_test.go

@@ -11,6 +11,7 @@ import (
 	"github.com/aws/aws-sdk-go/service/ec2"
 	"github.com/aws/aws-sdk-go/service/s3"
 	"github.com/mulgadc/spinifex/spinifex/awserrors"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/viperblock/viperblock"
@@ -99,7 +100,7 @@ func TestVolumeTagMirror_WritesOnlyControlPlaneTagsObject(t *testing.T) {
 	svc.natsConn = startTestNATS(t)
 	volumeID := "vol-enc-tag-writer"
 	seedEncryptedConfig(t, memoryStore, volumeID)
-	require.NoError(t, svc.putVolumeState(context.Background(), volumeID, volumeStateRecord{
+	require.NoError(t, svc.putVolumeState(context.Background(), volumeID, volumestate.Record{
 		State:            "in-use",
 		AttachedInstance: "i-live0000000000",
 		DeviceName:       "/dev/nbd0",

spinifex/handlers/ec2/volume/service_impl.go

@@ -22,6 +22,7 @@ import (
 	"github.com/mulgadc/spinifex/spinifex/config"
 	"github.com/mulgadc/spinifex/spinifex/filterutil"
 	handlers_ec2_instance "github.com/mulgadc/spinifex/spinifex/handlers/ec2/instance"
+	"github.com/mulgadc/spinifex/spinifex/handlers/ec2/volumestate"
 	"github.com/mulgadc/spinifex/spinifex/objectstore"
 	"github.com/mulgadc/spinifex/spinifex/types"
 	"github.com/mulgadc/spinifex/spinifex/utils"
@@ -1030,67 +1031,19 @@ type volumeConfigWrapper struct {
 	VolumeConfig viperblock.VolumeConfig `json:"VolumeConfig"`
 }
 
-// volumeStateRecord is the control-plane-owned attachment state, persisted to a
-// per-volume state.json object kept out of config.json. config.json is rewritten
-// by the live nbdkit VB on every SaveState (clobbering any State the control
-// plane wrote there) and is a sealed object for encrypted volumes (a second
-// writer reuses the AES-GCM nonce). state.json is plaintext, viperblock never
-// touches it, so the control plane is its single writer.
-type volumeStateRecord struct {
-	State            string    `json:"state"`
-	AttachedInstance string    `json:"attachedInstance"`
-	DeviceName       string    `json:"deviceName"`
-	AttachedAt       time.Time `json:"attachedAt"`
-}
-
-// volumeStateKey is the S3 key for a volume's control-plane state object.
-func volumeStateKey(volumeID string) string { return volumeID + "/state.json" }
-
 // volumeTagsKey is the S3 key for a volume's control-plane tags object.
 func volumeTagsKey(volumeID string) string { return volumeID + "/tags.json" }
 
 // putVolumeState writes the control-plane attachment state to state.json.
-func (s *VolumeServiceImpl) putVolumeState(ctx context.Context, volumeID string, rec volumeStateRecord) error {
-	data, err := json.Marshal(rec)
-	if err != nil {
-		return fmt.Errorf("marshal volume state: %w", err)
-	}
-	_, err = s.store.PutObject(ctx, &s3.PutObjectInput{
-		Bucket: aws.String(s.bucketName),
-		Key:    aws.String(volumeStateKey(volumeID)),
-		Body:   bytes.NewReader(data),
-	})
-	if err != nil {
-		return fmt.Errorf("failed to write volume state to S3: %w", err)
-	}
-	return nil
+func (s *VolumeServiceImpl) putVolumeState(ctx context.Context, volumeID string, rec volumestate.Record) error {
+	return volumestate.Write(ctx, s.store, s.bucketName, volumeID, rec)
 }
 
 // getVolumeState reads state.json. found=false with a nil error means the object
 // is absent (a volume predating the state.json split), in which case the caller
 // falls back to the State embedded in config.json.
-func (s *VolumeServiceImpl) getVolumeState(ctx context.Context, volumeID string) (volumeStateRecord, bool, error) {
-	getResult, err := s.store.GetObject(ctx, &s3.GetObjectInput{
-		Bucket: aws.String(s.bucketName),
-		Key:    aws.String(volumeStateKey(volumeID)),
-	})
-	if err != nil {
-		if objectstore.IsNoSuchKeyError(err) {
-			return volumeStateRecord{}, false, nil
-		}
-		return volumeStateRecord{}, false, fmt.Errorf("failed to get volume state: %w", err)
-	}
-	defer getResult.Body.Close()
-
-	body, err := io.ReadAll(getResult.Body)
-	if err != nil {
-		return volumeStateRecord{}, false, fmt.Errorf("failed to read volume state body: %w", err)
-	}
-	var rec volumeStateRecord
-	if err := json.Unmarshal(body, &rec); err != nil {
-		return volumeStateRecord{}, false, fmt.Errorf("failed to unmarshal volume state: %w", err)
-	}
-	return rec, true, nil
+func (s *VolumeServiceImpl) getVolumeState(ctx context.Context, volumeID string) (volumestate.Record, bool, error) {
+	return volumestate.Read(ctx, s.store, s.bucketName, volumeID)
 }
 
 // putVolumeTags writes the control-plane-owned tag set to tags.json.
@@ -1387,7 +1340,7 @@ func (s *VolumeServiceImpl) UpdateVolumeState(volumeID, state, attachedInstance,
 		state = "available"
 	}
 
-	rec := volumeStateRecord{
+	rec := volumestate.Record{
 		State:            state,
 		AttachedInstance: attachedInstance,
 		DeviceName:       deviceName,

spinifex/handlers/ec2/volumestate/volumestate.go

@@ -0,0 +1,78 @@
+// Package volumestate reads and writes the control-plane-owned attachment
+// state for an EBS volume, persisted to a per-volume state.json object.
+//
+// It is kept out of config.json because config.json is rewritten by the live
+// nbdkit VB on every SaveState (clobbering any State the control plane wrote
+// there) and is a sealed object for encrypted volumes (a second writer reuses
+// the AES-GCM nonce). state.json is plaintext, viperblock never touches it, so
+// the control plane is its single writer. Readers outside the volume service
+// (the snapshot path) need the same record, which is why it lives here.
+package volumestate
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"time"
+
+	"github.com/aws/aws-sdk-go/aws"
+	"github.com/aws/aws-sdk-go/service/s3"
+	"github.com/mulgadc/spinifex/spinifex/objectstore"
+)
+
+// Record is the control-plane-owned attachment state of a volume.
+type Record struct {
+	State            string    `json:"state"`
+	AttachedInstance string    `json:"attachedInstance"`
+	DeviceName       string    `json:"deviceName"`
+	AttachedAt       time.Time `json:"attachedAt"`
+}
+
+// Key is the S3 key for a volume's control-plane state object.
+func Key(volumeID string) string { return volumeID + "/state.json" }
+
+// Write persists the attachment state, replacing any previous record.
+func Write(ctx context.Context, store objectstore.ObjectStore, bucket, volumeID string, rec Record) error {
+	data, err := json.Marshal(rec)
+	if err != nil {
+		return fmt.Errorf("marshal volume state: %w", err)
+	}
+	_, err = store.PutObject(ctx, &s3.PutObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(Key(volumeID)),
+		Body:   bytes.NewReader(data),
+	})
+	if err != nil {
+		return fmt.Errorf("failed to write volume state to S3: %w", err)
+	}
+	return nil
+}
+
+// Read returns the attachment state. found=false with a nil error means the
+// object is absent (a volume predating the state.json split), in which case the
+// caller falls back to the State embedded in config.json.
+func Read(ctx context.Context, store objectstore.ObjectStore, bucket, volumeID string) (Record, bool, error) {
+	getResult, err := store.GetObject(ctx, &s3.GetObjectInput{
+		Bucket: aws.String(bucket),
+		Key:    aws.String(Key(volumeID)),
+	})
+	if err != nil {
+		if objectstore.IsNoSuchKeyError(err) {
+			return Record{}, false, nil
+		}
+		return Record{}, false, fmt.Errorf("failed to get volume state: %w", err)
+	}
+	defer getResult.Body.Close()
+
+	body, err := io.ReadAll(getResult.Body)
+	if err != nil {
+		return Record{}, false, fmt.Errorf("failed to read volume state body: %w", err)
+	}
+	var rec Record
+	if err := json.Unmarshal(body, &rec); err != nil {
+		return Record{}, false, fmt.Errorf("failed to unmarshal volume state: %w", err)
+	}
+	return rec, true, nil
+}

spinifex/testutil/drainsocket.go

@@ -0,0 +1,83 @@
+package testutil
+
+import (
+	"net"
+	"os"
+	"path/filepath"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+// SocketTempDir returns a data dir short enough to hold a unix socket path.
+// t.TempDir() embeds the test name, which pushes the drain socket underneath it
+// past the 108-byte sun_path limit for anything but the shortest test names.
+func SocketTempDir(t *testing.T) string {
+	t.Helper()
+
+	//nolint:usetesting // t.TempDir() yields a path too long for the unix socket
+	// sun_path (108-byte limit); MkdirTemp("") keeps it short under /tmp.
+	dir, err := os.MkdirTemp("", "spx")
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = os.RemoveAll(dir) })
+
+	return dir
+}
+
+// StartDrainSocket stands in for the NBD plugin's drain socket, at the path the
+// plugin serves it on the node hosting volumeID ({dataDir}/viperblock/{volumeID}
+// /snapshot.sock). Every connection is answered with ack ("OK\n" for a drain
+// that reached S3, "ERR\n" for one that did not) and closed.
+func StartDrainSocket(t *testing.T, dataDir, volumeID, ack string) {
+	t.Helper()
+	serveDrainSocket(t, dataDir, volumeID, ack, nil, nil)
+}
+
+// StartBlockingDrainSocket stands in for a plugin whose flush is still running:
+// it accepts but does not ack until release is called. accepted closes on the
+// first connection, so a test can be sure the drain is in flight before
+// asserting on what happens alongside it.
+func StartBlockingDrainSocket(t *testing.T, dataDir, volumeID, ack string) (accepted <-chan struct{}, release func()) {
+	t.Helper()
+
+	acceptedCh := make(chan struct{})
+	releaseCh := make(chan struct{})
+	serveDrainSocket(t, dataDir, volumeID, ack, acceptedCh, releaseCh)
+
+	var once sync.Once
+	t.Cleanup(func() { once.Do(func() { close(releaseCh) }) })
+
+	return acceptedCh, func() { once.Do(func() { close(releaseCh) }) }
+}
+
+// serveDrainSocket answers every connection with ack, optionally closing
+// accepted on the first one and waiting for release before it replies.
+func serveDrainSocket(t *testing.T, dataDir, volumeID, ack string, accepted chan struct{}, release <-chan struct{}) {
+	t.Helper()
+
+	dir := filepath.Join(dataDir, "viperblock", volumeID)
+	require.NoError(t, os.MkdirAll(dir, 0o750))
+
+	ln, err := net.Listen("unix", filepath.Join(dir, "snapshot.sock"))
+	require.NoError(t, err)
+	t.Cleanup(func() { _ = ln.Close() })
+
+	go func() {
+		var once sync.Once
+		for {
+			conn, err := ln.Accept()
+			if err != nil {
+				return // listener closed by the cleanup above
+			}
+			if accepted != nil {
+				once.Do(func() { close(accepted) })
+			}
+			if release != nil {
+				<-release
+			}
+			_, _ = conn.Write([]byte(ack))
+			_ = conn.Close()
+		}
+	}()
+}

spinifex/types/ec2.go

@@ -7,6 +7,7 @@ type EC2InstanceCommand struct {
 	Attributes                EC2CommandAttributes       `json:"attributes"`
 	AttachVolumeData          *AttachVolumeData          `json:"attach_volume_data,omitempty"`
 	DetachVolumeData          *DetachVolumeData          `json:"detach_volume_data,omitempty"`
+	DrainVolumeData           *DrainVolumeData           `json:"drain_volume_data,omitempty"`
 	AttachENIData             *AttachENIData             `json:"attach_eni_data,omitempty"`
 	DetachENIData             *DetachENIData             `json:"detach_eni_data,omitempty"`
 	IamProfileAssociationData *IamProfileAssociationData `json:"iam_profile_association_data,omitempty"`
@@ -21,6 +22,7 @@ type EC2CommandAttributes struct {
 	StartInstance               bool `json:"start_instance"`
 	AttachVolume                bool `json:"attach_volume"`
 	DetachVolume                bool `json:"detach_volume"`
+	DrainVolume                 bool `json:"drain_volume,omitempty"`
 	RebootInstance              bool `json:"reboot_instance"`
 	AttachENI                   bool `json:"attach_eni"`
 	DetachENI                   bool `json:"detach_eni"`
@@ -43,6 +45,33 @@ type DetachVolumeData struct {
 	Force    bool   `json:"force,omitempty"`
 }
 
+// DrainVolumeData carries parameters for a drain-volume command, which flushes
+// the volume's in-flight writes to S3 before a snapshot reads them. The command
+// is addressed to the instance the volume is attached to because the drain
+// socket only exists on the node hosting it.
+type DrainVolumeData struct {
+	VolumeID string `json:"volume_id"`
+}
+
+const (
+	// DrainVolumeStatusDrained is the ack a node returns once the volume's
+	// in-flight writes have reached S3.
+	DrainVolumeStatusDrained = "drained"
+
+	// DrainVolumeStatusNotRunning is the ack a node returns when it still holds
+	// the instance but the VM is not running, so no process is writing to the
+	// volume and its sealed checkpoint is already current. An attachment record
+	// outlives the writer — stop deliberately leaves boot volumes attached — so
+	// the caller cannot tell this from the record alone.
+	DrainVolumeStatusNotRunning = "not-running"
+)
+
+// DrainVolumeResponse is the reply to a drain-volume command.
+type DrainVolumeResponse struct {
+	VolumeID string `json:"volume_id"`
+	Status   string `json:"status"`
+}
+
 // AttachENIData carries parameters for an attach-network-interface command.
 type AttachENIData struct {
 	NetworkInterfaceID string `json:"network_interface_id"`

spinifex/vm/lifecycle.go

@@ -166,6 +166,19 @@ func (m *Manager) launch(ctx context.Context, instance *VM) (err error) {
 		return nil
 	}
 
+	// Make the routing record durable before QEMU can write. If persistence
+	// fails, unmounting seals the untouched backend and keeps the launch from
+	// exposing a writer that snapshots would classify as available.
+	if stateErr := m.markAttachedVolumesInUse(instance); stateErr != nil {
+		if unmountErr := m.deps.VolumeMounter.Unmount(instance); unmountErr != nil {
+			return errors.Join(
+				fmt.Errorf("persist attached volume state: %w", stateErr),
+				fmt.Errorf("rollback mounted volumes: %w", unmountErr),
+			)
+		}
+		return fmt.Errorf("persist attached volume state: %w", stateErr)
+	}
+
 	_, qemuSpan := otel.Tracer(vmTracerName).Start(ctx, "vm.launch.start_qemu")
 	qemuErr := m.startQEMU(instance)
 	endSpanWithError(qemuSpan, qemuErr)
@@ -212,9 +225,6 @@ func (m *Manager) launch(ctx context.Context, instance *VM) (err error) {
 		}
 	}
 
-	// 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
 		// launch still succeeds. The instance is reachable via cluster
@@ -229,27 +239,28 @@ func (m *Manager) launch(ctx context.Context, instance *VM) (err error) {
 	return nil
 }
 
-// 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
-	}
+// markAttachedVolumesInUse persists routing state for every attached EBS
+// volume before a launch exposes them to QEMU. Callers must abort or fail
+// recovery when an update fails; logging and continuing would let snapshots
+// mistake a live writer for an available volume.
+func (m *Manager) markAttachedVolumesInUse(instance *VM) error {
 	instance.EBSRequests.Mu.Lock()
-	defer instance.EBSRequests.Mu.Unlock()
-	for _, ebsReq := range instance.EBSRequests.Requests {
+	requests := append([]types.EBSRequest(nil), instance.EBSRequests.Requests...)
+	instance.EBSRequests.Mu.Unlock()
+
+	for _, ebsReq := range requests {
 		// EFI pflash is not a KV-tracked EBS volume; skip it like the unmount gate.
 		if ebsReq.EFI {
 			continue
 		}
+		if m.deps.VolumeStateUpdater == nil {
+			return errors.New("volume state updater not wired")
+		}
 		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)
+			return fmt.Errorf("persist in-use state for volume %s: %w", ebsReq.Name, err)
 		}
 	}
+	return nil
 }
 
 const (

spinifex/vm/lifecycle_test.go

@@ -423,6 +423,41 @@ func TestRun_VolumeMounterError_Propagates(t *testing.T) {
 	assert.Equal(t, StatePending, m.Status(instance), "status must be unchanged on Mount failure")
 }
 
+// TestRun_VolumeStateFailureRollsBackBeforeQEMU verifies launch does not expose
+// a writer when authoritative attachment state cannot be persisted.
+func TestRun_VolumeStateFailureRollsBackBeforeQEMU(t *testing.T) {
+	t.Setenv("XDG_RUNTIME_DIR", t.TempDir())
+	stateErr := errors.New("predastore unavailable")
+	mounter := &fakeVolumeMounter{}
+	stateUpdater := &fakeVolumeStateUpdater{err: stateErr}
+	m := NewManagerWithDeps(Deps{
+		VolumeMounter:      mounter,
+		VolumeStateUpdater: stateUpdater,
+	})
+	instance := &VM{
+		ID:     "i-state-fail",
+		Status: StatePending,
+		EBSRequests: types.EBSRequests{Requests: []types.EBSRequest{{
+			Name: "vol-1", DeviceName: "/dev/sda1", Boot: true,
+		}}},
+	}
+	m.Insert(instance)
+
+	err := m.Run(t.Context(), instance)
+
+	require.Error(t, err)
+	assert.ErrorIs(t, err, stateErr)
+	assert.Equal(t, []string{instance.ID}, mounter.mounted)
+	assert.Equal(t, []string{instance.ID}, mounter.unmounted,
+		"failed attachment-state persistence must roll back mounted volumes")
+	calls := stateUpdater.snapshot()
+	require.Len(t, calls, 1)
+	assert.Equal(t, "in-use", calls[0].State)
+	assert.Equal(t, instance.ID, calls[0].InstanceID)
+	assert.NoFileExists(t, filepath.Join(os.Getenv("XDG_RUNTIME_DIR"), instance.ID+".pid"),
+		"startQEMU must not run after attachment-state persistence fails")
+}
+
 // TestRun_AlreadyRunningPID_ReturnsError covers the live-PID guard
 // (lifecycle.go:81-91): when the PID file refers to a live process, Run
 // must return an error before touching volumes or firing hooks.

spinifex/vm/manager_deps_test.go

@@ -205,9 +205,10 @@ var _ VolumeMounter = (*fakeVolumeMounter)(nil)
 // of the API-form name (e.g. /dev/sdf) shows up as the wrong stored
 // device on the recorded call.
 type fakeVolumeStateUpdater struct {
-	mu    sync.Mutex
-	calls []volumeStateUpdate
-	err   error
+	mu       sync.Mutex
+	calls    []volumeStateUpdate
+	err      error
+	onUpdate func(volumeStateUpdate)
 }
 
 type volumeStateUpdate struct {
@@ -218,15 +219,23 @@ type volumeStateUpdate struct {
 }
 
 func (f *fakeVolumeStateUpdater) UpdateVolumeState(volumeID, state, instanceID, attachmentDevice string) error {
-	f.mu.Lock()
-	defer f.mu.Unlock()
-	f.calls = append(f.calls, volumeStateUpdate{
+	update := volumeStateUpdate{
 		VolumeID:         volumeID,
 		State:            state,
 		InstanceID:       instanceID,
 		AttachmentDevice: attachmentDevice,
-	})
-	return f.err
+	}
+
+	f.mu.Lock()
+	f.calls = append(f.calls, update)
+	err := f.err
+	hook := f.onUpdate
+	f.mu.Unlock()
+
+	if hook != nil {
+		hook(update)
+	}
+	return err
 }
 
 func (f *fakeVolumeStateUpdater) snapshot() []volumeStateUpdate {

spinifex/vm/restore.go

@@ -330,12 +330,18 @@ func (m *Manager) reconnectInstance(instance *VM) error {
 		}
 	}
 
-	instance.Status = StateRunning
+	// Re-assert in-use state before advertising the surviving QEMU as running.
+	// A failed write leaves snapshot routing ambiguous, so recovery must fail
+	// closed rather than continue with a guest-writable available volume.
+	if err := m.markAttachedVolumesInUse(instance); err != nil {
+		if instance.QMPClient != nil && instance.QMPClient.Conn != nil {
+			_ = instance.QMPClient.Conn.Close()
+			instance.QMPClient = nil
+		}
+		return fmt.Errorf("persist attached volume state during reconnect: %w", err)
+	}
 
-	// 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)
+	instance.Status = StateRunning
 
 	if err := m.writeRunningState(); err != nil {
 		return fmt.Errorf("failed to persist reconnected instance state: %w", err)

spinifex/vm/volumes.go

@@ -30,16 +30,18 @@ const blockdevDelMaxAttempts = 20
 // margin; the ebs.unmount seal keeps its separate unmountSealTimeout.
 const detachAggregateTimeout = 3 * time.Minute
 
-// rollbackUnmount best-effort unmounts a volume while unwinding a failed
-// AttachVolume. The volume never took guest writes, so the unmount seal is a
-// no-op; a failure is logged and tolerated rather than masking the attach error.
-func (m *Manager) rollbackUnmount(req types.EBSRequest) {
+// rollbackUnmount unmounts a volume while unwinding a failed AttachVolume.
+// Callers keep the volume non-available when the seal fails so a later attach
+// cannot discard local state that may still be owned by the backend.
+func (m *Manager) rollbackUnmount(req types.EBSRequest) error {
 	if m.deps.VolumeMounter == nil {
-		return
+		return nil
 	}
 	if err := m.deps.VolumeMounter.UnmountOne(req); err != nil {
 		slog.Warn("AttachVolume: rollback unmount failed", "volume", req.Name, "err", err)
+		return err
 	}
+	return nil
 }
 
 // delIothreadBestEffort removes the per-volume iothread object on a failed
@@ -54,8 +56,25 @@ func (m *Manager) delIothreadBestEffort(ctx context.Context, instance *VM, iothr
 	}
 }
 
-// AttachVolume hot-plugs a volume via the QMP pipeline (mount → blockdev-add →
-// device_add). Partial state is rolled back on failure. If device is empty, the
+// rollbackHotAttach removes the QMP block backend before unmounting it. A
+// failed blockdev deletion must leave the NBD server mounted because QEMU still
+// owns it; false tells the caller to retain the fail-closed in-use state.
+func (m *Manager) rollbackHotAttach(ctx context.Context, instance *VM, req types.EBSRequest, nodeName, iothreadID string) bool {
+	if _, err := sendQMPCommand(ctx, instance.QMPClient, qmp.QMPCommand{
+		Execute:   "blockdev-del",
+		Arguments: map[string]any{"node-name": nodeName},
+	}, instance.ID); err != nil {
+		slog.ErrorContext(ctx, "AttachVolume: rollback blockdev-del failed, skipping EBS unmount",
+			"volumeId", req.Name, "err", err)
+		return false
+	}
+
+	m.delIothreadBestEffort(ctx, instance, iothreadID, req.Name)
+	return m.rollbackUnmount(req) == nil
+}
+
+// AttachVolume hot-plugs a volume via the pipeline mount → blockdev-add →
+// persist in-use → device_add. Partial state is rolled back on failure. If device is empty, the
 // next free /dev/sd[f-p] slot is allocated. Instance must be in StateRunning.
 // Returns the AWS-API device name (/dev/sd[f-p]) echoed in the API response.
 func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string) (string, error) {
@@ -96,7 +115,7 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 		// Empty-URI response leaves backend NBD state ambiguous; unmount
 		// defensively to avoid orphaning a half-started mount.
 		if errors.Is(err, ErrMountAmbiguous) {
-			m.rollbackUnmount(ebsRequest)
+			_ = m.rollbackUnmount(ebsRequest)
 		}
 		return "", fmt.Errorf("mount volume %s: %w", volumeID, err)
 	}
@@ -104,7 +123,7 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 	serverType, socketPath, nbdHost, nbdPort, err := utils.ParseNBDURI(ebsRequest.NBDURI)
 	if err != nil {
 		slog.ErrorContext(ctx, "AttachVolume: failed to parse NBDURI", "uri", ebsRequest.NBDURI, "err", err)
-		m.rollbackUnmount(ebsRequest)
+		_ = m.rollbackUnmount(ebsRequest)
 		return "", fmt.Errorf("parse NBDURI: %w", err)
 	}
 	serverArg := NBDServerOpts{Type: serverType, Path: socketPath, Host: nbdHost, Port: nbdPort}.QMPArg()
@@ -127,7 +146,7 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 	instance.EBSRequests.Mu.Unlock()
 	if hotplugPort == 0 {
 		slog.ErrorContext(ctx, "AttachVolume: EBS hot-plug port pool exhausted", "volumeId", volumeID)
-		m.rollbackUnmount(ebsRequest)
+		_ = m.rollbackUnmount(ebsRequest)
 		return "", ErrAttachmentLimitExceeded
 	}
 	ebsRequest.HotplugPort = hotplugPort
@@ -140,7 +159,7 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 		},
 	}, instance.ID); err != nil {
 		slog.ErrorContext(ctx, "AttachVolume: QMP object-add iothread failed", "volumeId", volumeID, "err", err)
-		m.rollbackUnmount(ebsRequest)
+		_ = m.rollbackUnmount(ebsRequest)
 		return "", fmt.Errorf("QMP object-add iothread: %w", err)
 	}
 
@@ -156,10 +175,23 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 	}, instance.ID); err != nil {
 		slog.ErrorContext(ctx, "AttachVolume: QMP blockdev-add failed", "volumeId", volumeID, "err", err)
 		m.delIothreadBestEffort(ctx, instance, iothreadID, volumeID)
-		m.rollbackUnmount(ebsRequest)
+		_ = m.rollbackUnmount(ebsRequest)
 		return "", fmt.Errorf("QMP blockdev-add: %w", err)
 	}
 
+	// Persist routing state before device_add makes the volume guest-writable.
+	// The API-form device name must be stored because DescribeVolumes attachment
+	// filters do not match the guest virtio path. A failed or ambiguous write
+	// aborts the attach and leaves any resulting in-use record fail-closed.
+	if m.deps.VolumeStateUpdater == nil {
+		m.rollbackHotAttach(ctx, instance, ebsRequest, nodeName, iothreadID)
+		return "", errors.New("volume state updater not wired")
+	}
+	if err := m.deps.VolumeStateUpdater.UpdateVolumeState(volumeID, "in-use", instance.ID, device); err != nil {
+		m.rollbackHotAttach(ctx, instance, ebsRequest, nodeName, iothreadID)
+		return "", fmt.Errorf("persist in-use state for volume %s: %w", volumeID, err)
+	}
+
 	// The virtio-blk-pci device must land on a free hot-plug PCIe root port
 	// (hotplug-ebs{N}); pcie.0 rejects hot-plug. The port is allocated from
 	// live QEMU state above, independent of the AWS device name.
@@ -177,15 +209,11 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 	}, instance.ID); err != nil {
 		slog.ErrorContext(ctx, "AttachVolume: QMP device_add failed, rolling back blockdev",
 			"volumeId", volumeID, "err", err)
-		if _, delErr := sendQMPCommand(ctx, instance.QMPClient, qmp.QMPCommand{
-			Execute:   "blockdev-del",
-			Arguments: map[string]any{"node-name": nodeName},
-		}, instance.ID); delErr != nil {
-			slog.ErrorContext(ctx, "AttachVolume: rollback blockdev-del failed, skipping EBS unmount",
-				"volumeId", volumeID, "err", delErr)
-		} else {
-			m.delIothreadBestEffort(ctx, instance, iothreadID, volumeID)
-			m.rollbackUnmount(ebsRequest)
+		if m.rollbackHotAttach(ctx, instance, ebsRequest, nodeName, iothreadID) {
+			if stateErr := m.deps.VolumeStateUpdater.UpdateVolumeState(volumeID, "available", "", ""); stateErr != nil {
+				slog.ErrorContext(ctx, "AttachVolume: failed to restore available state after rollback",
+					"volumeId", volumeID, "err", stateErr)
+			}
 		}
 		return "", fmt.Errorf("QMP device_add: %w", err)
 	}
@@ -243,21 +271,6 @@ func (m *Manager) AttachVolume(ctx context.Context, id, volumeID, device string)
 		v.Instance.BlockDeviceMappings = append(v.Instance.BlockDeviceMappings, mapping)
 	})
 
-	// Volume metadata, by contrast, drives Volume.Attachments[].Device
-	// and the attachment.device filter on DescribeVolumes. The Terraform
-	// AWS provider polls that filter with the API-form name (/dev/sd[f-p])
-	// supplied in the .tf config, so storing the guest path here makes
-	// the filter reject every attached volume and the post-attach wait
-	// loop fails with "couldn't find resource". Always persist the API
-	// name even though it diverges from the BDM convention above —
-	// nothing in mulga-599 rewrites this field after attach.
-	if m.deps.VolumeStateUpdater != nil {
-		if err := m.deps.VolumeStateUpdater.UpdateVolumeState(volumeID, "in-use", instance.ID, device); err != nil {
-			slog.ErrorContext(ctx, "AttachVolume: failed to update volume metadata",
-				"volumeId", volumeID, "err", err)
-		}
-	}
-
 	if err := m.writeRunningState(); err != nil {
 		slog.ErrorContext(ctx, "AttachVolume: failed to write state", "err", err)
 	}

spinifex/vm/volumes_test.go

@@ -656,7 +656,10 @@ func TestReboot_QMPFailureSurfacesError(t *testing.T) {
 // wired with the supplied QMP client. Used by AttachVolume rollback tests.
 func attachVolumeRunningInstance(t *testing.T, qmpClient *qmp.QMPClient, mounter VolumeMounter) (*Manager, *VM) {
 	t.Helper()
-	m := NewManagerWithDeps(Deps{VolumeMounter: mounter})
+	m := NewManagerWithDeps(Deps{
+		VolumeMounter:      mounter,
+		VolumeStateUpdater: &fakeVolumeStateUpdater{},
+	})
 	v := &VM{
 		ID:        "i-1",
 		Status:    StateRunning,
@@ -764,6 +767,67 @@ func TestAttachVolume_BlockdevAddFailure_TriggersUnmountOne(t *testing.T) {
 	assert.Equal(t, []string{"vol-1"}, mounter.unmountedOne)
 }
 
+// TestAttachVolume_PersistsStateBeforeDeviceAdd verifies the routing record is
+// durable before QEMU exposes the volume to the guest.
+func TestAttachVolume_PersistsStateBeforeDeviceAdd(t *testing.T) {
+	recorder := &qmpRecorder{}
+	qmpClient, cancel := newMockQMPClient(t, func(cmd qmp.QMPCommand) map[string]any {
+		recorder.record(cmd)
+		if cmd.Execute == "query-block" {
+			return map[string]any{"return": []qmp.BlockDevice{{
+				Inserted: &qmp.BlockInserted{},
+				QDev:     "/machine/peripheral/vdisk-vol-1/hotplug-ebs1/virtio-backend",
+			}}}
+		}
+		return nil
+	})
+	defer cancel()
+
+	var commandsAtUpdate []string
+	stateUpdater := &fakeVolumeStateUpdater{}
+	stateUpdater.onUpdate = func(update volumeStateUpdate) {
+		if update.State == "in-use" {
+			commandsAtUpdate = recorder.executes()
+		}
+	}
+	mounter := &fakeVolumeMounter{mountOneURI: "nbd:unix:/tmp/test.sock"}
+	m := NewManagerWithDeps(Deps{VolumeMounter: mounter, VolumeStateUpdater: stateUpdater})
+	m.Insert(&VM{ID: "i-1", Status: StateRunning, Instance: &ec2.Instance{}, QMPClient: qmpClient})
+
+	_, err := m.AttachVolume(t.Context(), "i-1", "vol-1", "/dev/sdf")
+	require.NoError(t, err)
+	assert.Equal(t, []string{"object-add", "blockdev-add"}, commandsAtUpdate,
+		"in-use state must be durable before device_add makes the volume writable")
+	assert.Contains(t, recorder.executes(), "device_add")
+}
+
+// TestAttachVolume_StateUpdateFailurePreventsDeviceAdd verifies attachment
+// state is no longer best-effort: a failed write rolls back the backend before
+// the guest can access it.
+func TestAttachVolume_StateUpdateFailurePreventsDeviceAdd(t *testing.T) {
+	recorder := &qmpRecorder{}
+	qmpClient, cancel := newMockQMPClient(t, func(cmd qmp.QMPCommand) map[string]any {
+		recorder.record(cmd)
+		return nil
+	})
+	defer cancel()
+
+	stateErr := errors.New("predastore unavailable")
+	stateUpdater := &fakeVolumeStateUpdater{err: stateErr}
+	mounter := &fakeVolumeMounter{mountOneURI: "nbd:unix:/tmp/test.sock"}
+	m := NewManagerWithDeps(Deps{VolumeMounter: mounter, VolumeStateUpdater: stateUpdater})
+	m.Insert(&VM{ID: "i-1", Status: StateRunning, Instance: &ec2.Instance{}, QMPClient: qmpClient})
+
+	_, err := m.AttachVolume(t.Context(), "i-1", "vol-1", "/dev/sdf")
+	require.Error(t, err)
+	assert.ErrorIs(t, err, stateErr)
+	assert.Equal(t, []string{"object-add", "blockdev-add", "blockdev-del", "object-del"}, recorder.executes())
+	assert.NotContains(t, recorder.executes(), "device_add")
+	assert.Equal(t, []string{"vol-1"}, mounter.unmountedOne)
+	require.Len(t, stateUpdater.snapshot(), 1)
+	assert.Equal(t, "in-use", stateUpdater.snapshot()[0].State)
+}
+
 // TestAttachVolume_DeviceAddFailure_BlockdevDelOK_Unmounts covers the
 // device_add-fail-then-rollback path where blockdev-del succeeds. The
 // rollback chain must run blockdev-del then UnmountOne in order.
@@ -782,7 +846,9 @@ func TestAttachVolume_DeviceAddFailure_BlockdevDelOK_Unmounts(t *testing.T) {
 	defer cancel()
 
 	mounter := &fakeVolumeMounter{mountOneURI: "nbd:unix:/tmp/test.sock"}
-	m, _ := attachVolumeRunningInstance(t, qmpClient, mounter)
+	stateUpdater := &fakeVolumeStateUpdater{}
+	m := NewManagerWithDeps(Deps{VolumeMounter: mounter, VolumeStateUpdater: stateUpdater})
+	m.Insert(&VM{ID: "i-1", Status: StateRunning, Instance: &ec2.Instance{}, QMPClient: qmpClient})
 
 	_, err := m.AttachVolume(t.Context(), "i-1", "vol-1", "/dev/sdf")
 	require.Error(t, err)
@@ -790,6 +856,11 @@ func TestAttachVolume_DeviceAddFailure_BlockdevDelOK_Unmounts(t *testing.T) {
 	assert.Equal(t, []string{"object-add", "blockdev-add", "device_add", "blockdev-del", "object-del"}, recorder.executes())
 	assert.Equal(t, []string{"vol-1"}, mounter.unmountedOne,
 		"successful blockdev-del rollback must be followed by UnmountOne")
+	calls := stateUpdater.snapshot()
+	require.Len(t, calls, 2)
+	assert.Equal(t, "in-use", calls[0].State)
+	assert.Equal(t, "available", calls[1].State,
+		"state returns to available only after the backend was removed and sealed")
 }
 
 // TestAttachVolume_DeviceAddFailure_BlockdevDelFails_SkipsUnmountOne