feat: record a launch before replying to the caller (#994)

* feat: time the local and KV halves of a state write separately

RunInstances responds before vmMgr.Insert and WriteState, so the cost of
gating the respond on the local write is measurable today: it is the
ec2.RecordRunInstances span. Splitting persistState's timing by store says
which half that cost is, and the summed instance count says whether the
local write scales with the node's set size.

Duration is a sum plus an op count rather than a histogram, because ES|QL
cannot aggregate a native histogram field.

* fix: count only the records a KV state write actually touched

WriteRunningSet is digest-guarded and writes just the changed records, so
attributing the node's whole VM count to it overstated the per-record cost.
Use Written+Retired, and rename the metric to records.sum with a description
that spells out the two meanings: the whole set locally, the changed ones in
KV.

* feat: record a launch before replying to the caller

RunInstances replied with the reservation and only then inserted the records
and wrote state, so a caller holding an instance ID could describe it and get
nothing back, and a guest could read IMDS before this node's state file knew
the instance existed. Move the reply after the insert, the state write and the
per-instance subscribe loop, so by the time an ID reaches the customer the node
can answer for it and a terminate issued immediately finds a responder.

A local write that fails now fails the request: the file is what this node
restores from, so a launch it did not record is a launch nothing can recover.
The handler undoes the insert with DeleteIf and returns the capacity, rather
than running instances the caller was told had failed.

Measured on env19 first: the newly-serialised work is 7-23ms against a 2.3-10.8s
end-to-end launch, under 1%, so no split of the state write was needed.

spinifex/daemon/daemon.go

@@ -2613,7 +2613,9 @@ func (d *Daemon) persistState(nodeID string, vms map[string]*vm.VM) error {
 	// the local failure. Revision only advances when the local write (the
 	// source of truth Revision() describes) actually landed.
 	path := d.localStatePath()
+	localStart := time.Now()
 	localErr := WriteLocalStateBytes(path, localData)
+	recordStateWrite("local", len(vms), time.Since(localStart))
 	if localErr != nil {
 		slog.Error("Local state write failed", "path", path, "error", localErr)
 	} else {
@@ -2625,8 +2627,13 @@ func (d *Daemon) persistState(nodeID string, vms map[string]*vm.VM) error {
 	// a set that has not been written yet. vms is the caller's snapshot and
 	// cannot change underneath either write.
 	if d.jsManager != nil {
+		kvStart := time.Now()
 		d.jsManager.WriteNodeMarkerBestEffort(nodeID, kvSyncTimeout)
-		d.jsManager.WriteRunningSet(nodeID, d.config.AZ, vms)
+		res := d.jsManager.WriteRunningSet(nodeID, d.config.AZ, vms)
+		// The digest guard means only changed records are written, so the KV
+		// write's cost tracks that count, not the whole set the local write
+		// rewrites.
+		recordStateWrite("kv", res.Written+res.Retired, time.Since(kvStart))
 	}
 
 	if localErr != nil {

spinifex/daemon/daemon_handlers_instance.go

@@ -184,21 +184,41 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) string {
 		}
 		return outcomeError
 	}
-	if err := msg.Respond(jsonResponse); err != nil {
-		slog.Error("Failed to respond to NATS request", "err", err)
-	}
-
+	// Record before responding: a caller that holds an instance ID must be
+	// able to find it. IMDS reads this node's on-disk state, so a guest can
+	// query it the moment the reservation reaches the customer.
+	_, recordSpan := otel.Tracer(daemonTracerName).Start(ctx, "ec2.RecordRunInstances",
+		trace.WithAttributes(attribute.Int("instance.count", len(instances))))
 	for _, instance := range instances {
 		d.vmMgr.Insert(instance)
 	}
-	if err := d.WriteState(); err != nil {
-		slog.Error("handleEC2RunInstances failed to write initial state", "err", err)
+	writeErr := d.WriteState()
+	endOpSpan(recordSpan, writeErr)
+	if writeErr != nil {
+		// The local file is the source of truth this node restores from, so a
+		// launch it does not record is a launch nothing can recover. Fail the
+		// request and undo, rather than run instances the caller was told
+		// failed.
+		slog.Error("handleEC2RunInstances failed to write initial state", "err", writeErr)
+		for _, instance := range instances {
+			// DeleteIf, so a slot another path has already reclaimed is left
+			// alone. The failed write means the file never gained these
+			// entries, so removing them returns memory and disk to agreement.
+			d.vmMgr.DeleteIf(instance.ID, instance)
+			if reservationID == "" {
+				d.resourceMgr.deallocate(instanceType)
+			} else {
+				d.resourceMgr.ReleaseToReservation(reservationID, instanceType)
+			}
+		}
+		respondWithError(d.node, msg, awserrors.ErrorServerInternal)
+		return outcomeError
 	}
 	slog.Info("Instances added to state with pending status", "count", len(instances))
 
 	// Project launch tags into the central tag store so describe-tags agrees
-	// with the record from birth. Best-effort: the reservation is already
-	// returned, so a failed write is logged rather than failing the launch.
+	// with the record from birth. Best-effort: a failed write is logged rather
+	// than failing the launch.
 	for _, instance := range instances {
 		if instance.Instance == nil || len(instance.Instance.Tags) == 0 {
 			continue
@@ -210,8 +230,9 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) string {
 		}
 	}
 
-	// Subscribe per-instance NATS topics so terminate/stop reach this daemon
-	// while volumes prepare. LaunchInstance replaces these on success.
+	// Subscribe per-instance NATS topics before responding, so a terminate or
+	// stop issued the instant the caller has the ID reaches this daemon rather
+	// than finding no responder. LaunchInstance replaces these on success.
 	d.mu.Lock()
 	for _, instance := range instances {
 		sub, subErr := d.natsConn.Subscribe(fmt.Sprintf("ec2.cmd.%s", instance.ID), d.handleEC2Events)
@@ -223,6 +244,10 @@ func (d *Daemon) handleEC2RunInstances(msg *nats.Msg) string {
 	}
 	d.mu.Unlock()
 
+	if err := msg.Respond(jsonResponse); err != nil {
+		slog.Error("Failed to respond to NATS request", "err", err)
+	}
+
 	_, launchSpan := otel.Tracer(daemonTracerName).Start(ctx, "ec2.LaunchRunInstances",
 		trace.WithAttributes(attribute.Int("instance.count", len(instances))))
 	d.instanceService.LaunchRunInstances(ctx, instances, input, instanceType)

spinifex/daemon/daemon_handlers_instance_persist_test.go

@@ -0,0 +1,73 @@
+package daemon
+
+import (
+	"encoding/json"
+	"os"
+	"path/filepath"
+	"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/vm"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+// breakLocalStateWrite makes the daemon's local state write fail by occupying
+// its parent directory path with a regular file, so MkdirAll returns ENOTDIR.
+// Everything else rooted at DataDir keeps working.
+func breakLocalStateWrite(t *testing.T, d *Daemon) {
+	t.Helper()
+	stateDir := filepath.Dir(d.localStatePath())
+	require.NoError(t, os.RemoveAll(stateDir))
+	require.NoError(t, os.WriteFile(stateDir, []byte("not a directory"), 0o600))
+	require.Error(t, d.WriteState(), "state write must fail once the path is blocked")
+}
+
+func countVMs(d *Daemon) int {
+	n := 0
+	d.vmMgr.View(func(vms map[string]*vm.VM) { n = len(vms) })
+	return n
+}
+
+// A launch whose local state write fails is reported as a failure and leaves
+// nothing behind: no record in the manager, and the capacity it reserved is
+// returned. The customer is never told "no" about an instance that exists.
+func TestHandleEC2RunInstances_LocalWriteFailureLaunchesNothing(t *testing.T) {
+	daemon, memStore := createFullTestDaemonWithStore(t, sharedNATSURL)
+	seedTestAMI(t, memStore, daemon.config.Predastore.Bucket, "ami-writefail")
+
+	instanceType := getTestInstanceType(t)
+	typeInfo := daemon.resourceMgr.InstanceTypes()[instanceType]
+	require.NotNil(t, typeInfo)
+	capacityBefore := daemon.resourceMgr.CanAllocate(typeInfo, 64)
+
+	sub, err := daemon.natsConn.QueueSubscribe("ec2.RunInstances.writefail", "spinifex-workers", asMsgHandler(daemon.handleEC2RunInstances))
+	require.NoError(t, err)
+	defer func() { _ = sub.Unsubscribe() }()
+
+	breakLocalStateWrite(t, daemon)
+
+	input := &ec2.RunInstancesInput{
+		ImageId:      aws.String("ami-writefail"),
+		InstanceType: aws.String(instanceType),
+		MinCount:     aws.Int64(1),
+		MaxCount:     aws.Int64(1),
+	}
+	reply, err := natsRequest(daemon.natsConn, "ec2.RunInstances.writefail", mustMarshal(t, input), 5*time.Second)
+	require.NoError(t, err)
+
+	assert.Equal(t, awserrors.ErrorServerInternal, decodeError(t, reply.Data)["Code"],
+		"a launch that could not be recorded must be reported as an error: %s", reply.Data)
+
+	var reservation ec2.Reservation
+	if json.Unmarshal(reply.Data, &reservation) == nil {
+		assert.Empty(t, reservation.Instances, "no instance may be returned to the caller")
+	}
+
+	assert.Zero(t, countVMs(daemon), "the failed launch must leave no record behind")
+	assert.Equal(t, capacityBefore, daemon.resourceMgr.CanAllocate(typeInfo, 64),
+		"capacity reserved for the failed launch must be returned")
+}

spinifex/daemon/daemon_handlers_instance_tags_test.go

@@ -153,7 +153,7 @@ func TestHandleSetInstanceTags_CrossAccountRejected(t *testing.T) {
 
 // RunInstances with instance-scoped TagSpecifications projects the launch tags
 // into the central tag store, so describe-tags sees them from birth. The write
-// happens after the reservation reply, hence the Eventually.
+// is best-effort, so the assertion tolerates it landing slightly late.
 func TestHandleEC2RunInstances_LaunchTagsWriteCentralStore(t *testing.T) {
 	daemon, memStore := createFullTestDaemonWithStore(t, sharedNATSURL)
 	seedTestAMI(t, memStore, daemon.config.Predastore.Bucket, "ami-launchtags")

spinifex/daemon/state_write_metrics.go

@@ -0,0 +1,74 @@
+package daemon
+
+import (
+	"context"
+	"sync"
+	"time"
+
+	"go.opentelemetry.io/otel"
+	"go.opentelemetry.io/otel/attribute"
+	"go.opentelemetry.io/otel/metric"
+)
+
+// stateWriteMeterName is the instrumentation scope for persistState timings.
+const stateWriteMeterName = "github.com/mulgadc/spinifex/spinifex/daemon/statewrite"
+
+// Duration is reported as a sum plus an op count rather than a histogram:
+// ES|QL cannot aggregate a native histogram field, so a panel computes
+// avg = SUM(duration.sum) / SUM(ops).
+type stateWriteInstruments struct {
+	ops      metric.Int64Counter
+	duration metric.Float64Counter
+	vms      metric.Int64Counter
+}
+
+var (
+	stateWriteOnce sync.Once
+	stateWriteInst stateWriteInstruments
+)
+
+func stateWriteMetrics() stateWriteInstruments {
+	stateWriteOnce.Do(func() {
+		meter := otel.Meter(stateWriteMeterName)
+
+		var err error
+		stateWriteInst.ops, err = meter.Int64Counter("mulga.daemon.state_write.ops",
+			metric.WithDescription("Count of local-file and KV state writes, by store."),
+			metric.WithUnit("{write}"))
+		if err != nil {
+			otel.Handle(err)
+		}
+		stateWriteInst.duration, err = meter.Float64Counter("mulga.daemon.state_write.duration.sum",
+			metric.WithDescription("Summed wall time of state writes, by store. Divide by ops for the mean."),
+			metric.WithUnit("s"))
+		if err != nil {
+			otel.Handle(err)
+		}
+		stateWriteInst.vms, err = meter.Int64Counter("mulga.daemon.state_write.records.sum",
+			metric.WithDescription("Summed records each write actually handled: the whole set for store=local, only the changed ones for store=kv. Divide by ops for the mean."),
+			metric.WithUnit("{record}"))
+		if err != nil {
+			otel.Handle(err)
+		}
+	})
+	return stateWriteInst
+}
+
+// recordStateWrite reports one write of store ("local" or "kv") over records,
+// which is the whole set locally but only the changed ones in KV. Every
+// recorder tolerates a nil instrument.
+func recordStateWrite(store string, records int, d time.Duration) {
+	inst := stateWriteMetrics()
+	set := metric.WithAttributeSet(attribute.NewSet(attribute.String("store", store)))
+	ctx := context.Background()
+
+	if inst.ops != nil {
+		inst.ops.Add(ctx, 1, set)
+	}
+	if inst.duration != nil {
+		inst.duration.Add(ctx, d.Seconds(), set)
+	}
+	if inst.vms != nil {
+		inst.vms.Add(ctx, int64(records), set)
+	}
+}