feat: O(1) snapshot volume checks (#38)

hive/daemon/daemon.go

@@ -592,8 +592,12 @@ func (d *Daemon) Start() error {
 	d.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(d.config, d.resourceMgr.instanceTypes, d.natsConn, &d.Instances, store)
 	d.keyService = handlers_ec2_key.NewKeyServiceImpl(d.config)
 	d.imageService = handlers_ec2_image.NewImageServiceImpl(d.config, d.natsConn)
-	d.volumeService = handlers_ec2_volume.NewVolumeServiceImpl(d.config, d.natsConn)
-	d.snapshotService = handlers_ec2_snapshot.NewSnapshotServiceImpl(d.config, d.natsConn)
+	snapSvc, snapshotKV, err := handlers_ec2_snapshot.NewSnapshotServiceImplWithNATS(d.config, d.natsConn)
+	if err != nil {
+		return fmt.Errorf("failed to initialize snapshot service with NATS KV: %w", err)
+	}
+	d.snapshotService = snapSvc
+	d.volumeService = handlers_ec2_volume.NewVolumeServiceImpl(d.config, d.natsConn, snapshotKV)
 	d.tagsService = handlers_ec2_tags.NewTagsServiceImpl(d.config)
 
 	if eigwSvc, err := handlers_ec2_eigw.NewEgressOnlyIGWServiceImplWithNATS(d.config, d.natsConn); err != nil {

hive/daemon/daemon_test.go

@@ -76,7 +76,7 @@ func createTestDaemon(t *testing.T, natsURL string) *Daemon {
 
 	// Initialize services (needed for handler tests)
 	daemon.instanceService = handlers_ec2_instance.NewInstanceServiceImpl(cfg, daemon.resourceMgr.instanceTypes, nc, &daemon.Instances, objectstore.NewMemoryObjectStore())
-	daemon.volumeService = handlers_ec2_volume.NewVolumeServiceImpl(cfg, nc)
+	daemon.volumeService = handlers_ec2_volume.NewVolumeServiceImpl(cfg, nc, nil)
 
 	t.Cleanup(func() {
 		if daemon.natsConn != nil {

hive/handlers/ec2/snapshot/service_impl.go

@@ -24,11 +24,14 @@ import (
 // Ensure SnapshotServiceImpl implements SnapshotService
 var _ SnapshotService = (*SnapshotServiceImpl)(nil)
 
+const KVBucketVolumeSnapshots = "hive-volume-snapshots"
+
 // SnapshotServiceImpl implements SnapshotService with S3-backed storage
 type SnapshotServiceImpl struct {
 	config   *config.Config
 	store    objectstore.ObjectStore
 	natsConn *nats.Conn
+	snapKV   nats.KeyValue
 	mutex    sync.RWMutex
 }
 
@@ -47,8 +50,8 @@ type SnapshotConfig struct {
 	Tags             map[string]string `json:"tags"`
 }
 
-// NewSnapshotServiceImpl creates a new snapshot service implementation
-func NewSnapshotServiceImpl(cfg *config.Config, natsConn *nats.Conn) *SnapshotServiceImpl {
+// NewSnapshotServiceImplWithNATS creates a snapshot service with JetStream KV for volume-snapshot tracking
+func NewSnapshotServiceImplWithNATS(cfg *config.Config, natsConn *nats.Conn) (*SnapshotServiceImpl, nats.KeyValue, error) {
 	store := objectstore.NewS3ObjectStoreFromConfig(
 		cfg.Predastore.Host,
 		cfg.Predastore.Region,
@@ -56,20 +59,52 @@ func NewSnapshotServiceImpl(cfg *config.Config, natsConn *nats.Conn) *SnapshotSe
 		cfg.Predastore.SecretKey,
 	)
 
+	js, err := natsConn.JetStream()
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to get JetStream context: %w", err)
+	}
+
+	kv, err := getOrCreateKVBucket(js, KVBucketVolumeSnapshots, 10)
+	if err != nil {
+		return nil, nil, fmt.Errorf("failed to create KV bucket %s: %w", KVBucketVolumeSnapshots, err)
+	}
+
+	slog.Info("Snapshot service initialized with JetStream KV", "bucket", KVBucketVolumeSnapshots)
+
 	return &SnapshotServiceImpl{
 		config:   cfg,
 		store:    store,
 		natsConn: natsConn,
-	}
+		snapKV:   kv,
+	}, kv, nil
 }
 
-// NewSnapshotServiceImplWithStore creates a snapshot service with a custom ObjectStore (for testing)
-func NewSnapshotServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn) *SnapshotServiceImpl {
-	return &SnapshotServiceImpl{
+// NewSnapshotServiceImplWithStore creates a snapshot service with a custom ObjectStore (for testing).
+// An optional snapshotKV can be provided for KV-backed volume-snapshot tracking.
+func NewSnapshotServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn, snapshotKV ...nats.KeyValue) *SnapshotServiceImpl {
+	svc := &SnapshotServiceImpl{
 		config:   cfg,
 		store:    store,
 		natsConn: natsConn,
 	}
+	if len(snapshotKV) > 0 {
+		svc.snapKV = snapshotKV[0]
+	}
+	return svc
+}
+
+func getOrCreateKVBucket(js nats.JetStreamContext, bucketName string, history int) (nats.KeyValue, error) {
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+		Bucket:  bucketName,
+		History: utils.SafeIntToUint8(history),
+	})
+	if err != nil {
+		kv, err = js.KeyValue(bucketName)
+		if err != nil {
+			return nil, err
+		}
+	}
+	return kv, nil
 }
 
 // getSnapshotKey returns the S3 key for storing snapshot metadata.
@@ -197,22 +232,25 @@ func (s *SnapshotServiceImpl) CreateSnapshot(input *ec2.CreateSnapshotInput) (*e
 		}
 
 		msg, err := s.natsConn.Request(fmt.Sprintf("ebs.snapshot.%s", volumeID), snapData, 30*time.Second)
-		if err != nil {
+		if errors.Is(err, nats.ErrNoResponders) {
+			// Volume is not mounted — data is already persisted to S3, proceed with metadata-only snapshot.
+			slog.Info("CreateSnapshot: volume not mounted, creating metadata-only snapshot", "volumeId", volumeID, "snapshotId", snapshotID)
+		} else if err != nil {
 			slog.Error("CreateSnapshot: ebs.snapshot NATS request failed", "volumeId", volumeID, "snapshotId", snapshotID, "err", err)
 			return nil, errors.New(awserrors.ErrorServerInternal)
-		}
+		} else {
+			var snapResp config.EBSSnapshotResponse
+			if err := json.Unmarshal(msg.Data, &snapResp); err != nil {
+				slog.Error("CreateSnapshot: failed to unmarshal ebs.snapshot response", "volumeId", volumeID, "snapshotId", snapshotID, "err", err)
+				return nil, errors.New(awserrors.ErrorServerInternal)
+			}
+			if !snapResp.Success || snapResp.Error != "" {
+				slog.Error("CreateSnapshot: viperblock snapshot failed", "volumeId", volumeID, "snapshotId", snapshotID, "err", snapResp.Error)
+				return nil, errors.New(awserrors.ErrorServerInternal)
+			}
 
-		var snapResp config.EBSSnapshotResponse
-		if err := json.Unmarshal(msg.Data, &snapResp); err != nil {
-			slog.Error("CreateSnapshot: failed to unmarshal ebs.snapshot response", "volumeId", volumeID, "snapshotId", snapshotID, "err", err)
-			return nil, errors.New(awserrors.ErrorServerInternal)
+			slog.Info("CreateSnapshot: viperblock snapshot created", "volumeId", volumeID, "snapshotId", snapshotID)
 		}
-		if !snapResp.Success || snapResp.Error != "" {
-			slog.Error("CreateSnapshot: viperblock snapshot failed", "volumeId", volumeID, "snapshotId", snapshotID, "err", snapResp.Error)
-			return nil, errors.New(awserrors.ErrorServerInternal)
-		}
-
-		slog.Info("CreateSnapshot: viperblock snapshot created", "volumeId", volumeID, "snapshotId", snapshotID)
 	} else {
 		slog.Warn("CreateSnapshot: natsConn is nil, skipping viperblock snapshot (metadata-only)", "volumeId", volumeID)
 	}
@@ -246,8 +284,16 @@ func (s *SnapshotServiceImpl) CreateSnapshot(input *ec2.CreateSnapshotInput) (*e
 		}
 	}
 
+	// Track the volume→snapshot dependency in KV before persisting to S3.
+	// This ensures we never have an untracked snapshot in S3.
+	if err := s.addSnapshotRef(volumeID, snapshotID); err != nil {
+		slog.Error("CreateSnapshot failed to add snapshot ref to KV", "snapshotId", snapshotID, "volumeId", volumeID, "err", err)
+		return nil, errors.New(awserrors.ErrorServerInternal)
+	}
+
 	if err := s.putSnapshotConfig(snapshotID, snapshotCfg); err != nil {
 		slog.Error("CreateSnapshot failed to write config", "snapshotId", snapshotID, "err", err)
+		_ = s.removeSnapshotRef(volumeID, snapshotID) // best-effort cleanup
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -318,7 +364,7 @@ func (s *SnapshotServiceImpl) DeleteSnapshot(input *ec2.DeleteSnapshotInput) (*e
 
 	slog.Info("DeleteSnapshot request", "snapshotId", snapshotID)
 
-	_, err := s.getSnapshotConfig(snapshotID)
+	cfg, err := s.getSnapshotConfig(snapshotID)
 	if err != nil {
 		slog.Error("DeleteSnapshot snapshot not found", "snapshotId", snapshotID, "err", err)
 		return nil, err
@@ -346,6 +392,12 @@ func (s *SnapshotServiceImpl) DeleteSnapshot(input *ec2.DeleteSnapshotInput) (*e
 		}
 	}
 
+	// Remove from KV after S3 cleanup. Failure is logged but not fatal —
+	// a phantom entry safely blocks volume deletion rather than allowing it.
+	if err := s.removeSnapshotRef(cfg.VolumeID, snapshotID); err != nil {
+		slog.Warn("DeleteSnapshot failed to remove snapshot ref from KV", "snapshotId", snapshotID, "volumeId", cfg.VolumeID, "err", err)
+	}
+
 	slog.Info("DeleteSnapshot completed", "snapshotId", snapshotID)
 
 	return &ec2.DeleteSnapshotOutput{}, nil
@@ -391,8 +443,15 @@ func (s *SnapshotServiceImpl) CopySnapshot(input *ec2.CopySnapshotInput) (*ec2.C
 		newCfg.Tags[k] = v
 	}
 
+	// Track the volume→snapshot dependency in KV before persisting to S3.
+	if err := s.addSnapshotRef(sourceCfg.VolumeID, newSnapshotID); err != nil {
+		slog.Error("CopySnapshot failed to add snapshot ref to KV", "snapshotId", newSnapshotID, "volumeId", sourceCfg.VolumeID, "err", err)
+		return nil, errors.New(awserrors.ErrorServerInternal)
+	}
+
 	if err := s.putSnapshotConfig(newSnapshotID, newCfg); err != nil {
 		slog.Error("CopySnapshot failed to write config", "snapshotId", newSnapshotID, "err", err)
+		_ = s.removeSnapshotRef(sourceCfg.VolumeID, newSnapshotID) // best-effort cleanup
 		return nil, errors.New(awserrors.ErrorServerInternal)
 	}
 
@@ -402,3 +461,137 @@ func (s *SnapshotServiceImpl) CopySnapshot(input *ec2.CopySnapshotInput) (*ec2.C
 		SnapshotId: aws.String(newSnapshotID),
 	}, nil
 }
+
+// addSnapshotRef adds snapshotID to the volume's snapshot list in KV.
+// Uses CAS (Create/Update with revision) to prevent lost updates under concurrency.
+func (s *SnapshotServiceImpl) addSnapshotRef(volumeID, snapshotID string) error {
+	if s.snapKV == nil {
+		slog.Debug("addSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID)
+		return nil
+	}
+
+	const maxRetries = 5
+	for attempt := range maxRetries {
+		entry, err := s.snapKV.Get(volumeID)
+		var snapshots []string
+
+		if err != nil {
+			if !errors.Is(err, nats.ErrKeyNotFound) {
+				return fmt.Errorf("addSnapshotRef: failed to get KV key %s: %w", volumeID, err)
+			}
+			// Key doesn't exist yet — create with just this snapshot
+			data, err := json.Marshal([]string{snapshotID})
+			if err != nil {
+				return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err)
+			}
+			if _, err := s.snapKV.Create(volumeID, data); err != nil {
+				if attempt < maxRetries-1 {
+					continue // concurrent Create/Update — retry
+				}
+				return fmt.Errorf("addSnapshotRef: failed to create KV key %s: %w", volumeID, err)
+			}
+			slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
+			return nil
+		}
+
+		if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
+			return fmt.Errorf("addSnapshotRef: failed to unmarshal KV value for %s: %w", volumeID, err)
+		}
+
+		snapshots = append(snapshots, snapshotID)
+
+		data, err := json.Marshal(snapshots)
+		if err != nil {
+			return fmt.Errorf("addSnapshotRef: failed to marshal snapshot list: %w", err)
+		}
+
+		if _, err := s.snapKV.Update(volumeID, data, entry.Revision()); err != nil {
+			if attempt < maxRetries-1 {
+				continue // concurrent update — retry
+			}
+			return fmt.Errorf("addSnapshotRef: failed to update KV key %s: %w", volumeID, err)
+		}
+
+		slog.Info("addSnapshotRef: added snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
+		return nil
+	}
+
+	return fmt.Errorf("addSnapshotRef: exhausted retries for KV key %s", volumeID)
+}
+
+// removeSnapshotRef removes snapshotID from the volume's snapshot list in KV.
+// Deletes the key if the list becomes empty.
+// Uses CAS (Update with revision) to prevent lost updates under concurrency.
+func (s *SnapshotServiceImpl) removeSnapshotRef(volumeID, snapshotID string) error {
+	if s.snapKV == nil {
+		slog.Debug("removeSnapshotRef: snapshotKV is nil, skipping", "volumeId", volumeID, "snapshotId", snapshotID)
+		return nil
+	}
+
+	const maxRetries = 5
+	for attempt := range maxRetries {
+		entry, err := s.snapKV.Get(volumeID)
+		if err != nil {
+			if errors.Is(err, nats.ErrKeyNotFound) {
+				return nil
+			}
+			return fmt.Errorf("removeSnapshotRef: failed to get KV key %s: %w", volumeID, err)
+		}
+
+		var snapshots []string
+		if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
+			return fmt.Errorf("removeSnapshotRef: failed to unmarshal KV value for %s: %w", volumeID, err)
+		}
+
+		filtered := snapshots[:0]
+		for _, snap := range snapshots {
+			if snap != snapshotID {
+				filtered = append(filtered, snap)
+			}
+		}
+
+		if len(filtered) == 0 {
+			if err := s.snapKV.Delete(volumeID); err != nil && !errors.Is(err, nats.ErrKeyNotFound) {
+				return fmt.Errorf("removeSnapshotRef: failed to delete KV key %s: %w", volumeID, err)
+			}
+		} else {
+			data, err := json.Marshal(filtered)
+			if err != nil {
+				return fmt.Errorf("removeSnapshotRef: failed to marshal snapshot list: %w", err)
+			}
+			if _, err := s.snapKV.Update(volumeID, data, entry.Revision()); err != nil {
+				if attempt < maxRetries-1 {
+					continue // concurrent update — retry
+				}
+				return fmt.Errorf("removeSnapshotRef: failed to update KV key %s: %w", volumeID, err)
+			}
+		}
+
+		slog.Info("removeSnapshotRef: removed snapshot ref", "volumeId", volumeID, "snapshotId", snapshotID)
+		return nil
+	}
+
+	return fmt.Errorf("removeSnapshotRef: exhausted retries for KV key %s", volumeID)
+}
+
+// volumeHasSnapshots returns true if the volume has any snapshots in KV.
+func (s *SnapshotServiceImpl) volumeHasSnapshots(volumeID string) (bool, error) {
+	if s.snapKV == nil {
+		return false, nil
+	}
+
+	entry, err := s.snapKV.Get(volumeID)
+	if err != nil {
+		if errors.Is(err, nats.ErrKeyNotFound) {
+			return false, nil
+		}
+		return false, fmt.Errorf("volumeHasSnapshots: failed to get KV key %s: %w", volumeID, err)
+	}
+
+	var snapshots []string
+	if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
+		return false, fmt.Errorf("volumeHasSnapshots: failed to unmarshal KV value for %s: %w", volumeID, err)
+	}
+
+	return len(snapshots) > 0, nil
+}

hive/handlers/ec2/snapshot/service_impl_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -12,6 +13,8 @@ import (
 	"github.com/mulgadc/hive/hive/config"
 	"github.com/mulgadc/hive/hive/objectstore"
 	"github.com/mulgadc/viperblock/viperblock"
+	"github.com/nats-io/nats-server/v2/server"
+	"github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -349,3 +352,209 @@ func TestCopySnapshot_PreservesTags(t *testing.T) {
 	assert.Equal(t, "Environment", *result.Snapshots[0].Tags[0].Key)
 	assert.Equal(t, "test", *result.Snapshots[0].Tags[0].Value)
 }
+
+// setupTestNATSKV creates a NATS JetStream test server and returns a KV bucket for testing.
+func setupTestNATSKV(t *testing.T) nats.KeyValue {
+	t.Helper()
+	opts := &server.Options{
+		Host:      "127.0.0.1",
+		Port:      -1,
+		JetStream: true,
+		StoreDir:  t.TempDir(),
+		NoLog:     true,
+		NoSigs:    true,
+	}
+	ns, err := server.NewServer(opts)
+	require.NoError(t, err)
+	go ns.Start()
+	require.True(t, ns.ReadyForConnections(5*time.Second))
+	t.Cleanup(func() { ns.Shutdown() })
+
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	t.Cleanup(func() { nc.Close() })
+
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+		Bucket: KVBucketVolumeSnapshots,
+	})
+	require.NoError(t, err)
+	return kv
+}
+
+func TestAddSnapshotRef(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := &SnapshotServiceImpl{snapKV: kv}
+
+	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
+	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-b"))
+
+	entry, err := kv.Get("vol-1")
+	require.NoError(t, err)
+	var snapshots []string
+	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
+	assert.Equal(t, []string{"snap-a", "snap-b"}, snapshots)
+}
+
+func TestRemoveSnapshotRef(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := &SnapshotServiceImpl{snapKV: kv}
+
+	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
+	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-b"))
+
+	// Remove one
+	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-a"))
+
+	entry, err := kv.Get("vol-1")
+	require.NoError(t, err)
+	var snapshots []string
+	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
+	assert.Equal(t, []string{"snap-b"}, snapshots)
+
+	// Remove last — key should be deleted
+	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-b"))
+
+	_, err = kv.Get("vol-1")
+	assert.ErrorIs(t, err, nats.ErrKeyNotFound)
+}
+
+func TestRemoveSnapshotRef_NonExistentKey(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := &SnapshotServiceImpl{snapKV: kv}
+
+	// Should not error on non-existent key
+	require.NoError(t, svc.removeSnapshotRef("vol-nonexistent", "snap-x"))
+}
+
+func TestVolumeHasSnapshots(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	svc := &SnapshotServiceImpl{snapKV: kv}
+
+	// No entry → false
+	has, err := svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, err)
+	assert.False(t, has)
+
+	// Add one → true
+	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
+	has, err = svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, err)
+	assert.True(t, has)
+
+	// Remove it → false
+	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-a"))
+	has, err = svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, err)
+	assert.False(t, has)
+}
+
+func TestKVNilFallback(t *testing.T) {
+	svc := &SnapshotServiceImpl{snapKV: nil}
+
+	// All methods should be no-ops when KV is nil
+	require.NoError(t, svc.addSnapshotRef("vol-1", "snap-a"))
+	require.NoError(t, svc.removeSnapshotRef("vol-1", "snap-a"))
+	has, err := svc.volumeHasSnapshots("vol-1")
+	require.NoError(t, err)
+	assert.False(t, has)
+}
+
+func TestCreateSnapshot_WritesKVEntry(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	store := objectstore.NewMemoryObjectStore()
+	cfg := &config.Config{
+		Predastore: config.PredastoreConfig{
+			Bucket:    "test-bucket",
+			AccessKey: "test-owner-123",
+		},
+	}
+	svc := NewSnapshotServiceImplWithStore(cfg, store, nil, kv)
+
+	volumeID := "vol-kvtest"
+	createTestVolume(t, store, volumeID, 10)
+
+	snap, err := svc.CreateSnapshot(&ec2.CreateSnapshotInput{
+		VolumeId: aws.String(volumeID),
+	})
+	require.NoError(t, err)
+
+	// Verify KV entry exists
+	has, err := svc.volumeHasSnapshots(volumeID)
+	require.NoError(t, err)
+	assert.True(t, has)
+
+	// Verify snapshot ID is in the list
+	entry, err := kv.Get(volumeID)
+	require.NoError(t, err)
+	var snapshots []string
+	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
+	assert.Contains(t, snapshots, *snap.SnapshotId)
+}
+
+func TestDeleteSnapshot_RemovesKVEntry(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	store := objectstore.NewMemoryObjectStore()
+	cfg := &config.Config{
+		Predastore: config.PredastoreConfig{
+			Bucket:    "test-bucket",
+			AccessKey: "test-owner-123",
+		},
+	}
+	svc := NewSnapshotServiceImplWithStore(cfg, store, nil, kv)
+
+	volumeID := "vol-kvdelete"
+	createTestVolume(t, store, volumeID, 10)
+
+	snap, err := svc.CreateSnapshot(&ec2.CreateSnapshotInput{
+		VolumeId: aws.String(volumeID),
+	})
+	require.NoError(t, err)
+
+	// Delete the snapshot
+	_, err = svc.DeleteSnapshot(&ec2.DeleteSnapshotInput{
+		SnapshotId: snap.SnapshotId,
+	})
+	require.NoError(t, err)
+
+	// KV should now be empty for this volume
+	has, err := svc.volumeHasSnapshots(volumeID)
+	require.NoError(t, err)
+	assert.False(t, has)
+}
+
+func TestCopySnapshot_AddsKVEntry(t *testing.T) {
+	kv := setupTestNATSKV(t)
+	store := objectstore.NewMemoryObjectStore()
+	cfg := &config.Config{
+		Predastore: config.PredastoreConfig{
+			Bucket:    "test-bucket",
+			AccessKey: "test-owner-123",
+		},
+	}
+	svc := NewSnapshotServiceImplWithStore(cfg, store, nil, kv)
+
+	volumeID := "vol-kvcopy"
+	createTestVolume(t, store, volumeID, 10)
+
+	snap, err := svc.CreateSnapshot(&ec2.CreateSnapshotInput{
+		VolumeId: aws.String(volumeID),
+	})
+	require.NoError(t, err)
+
+	copyResult, err := svc.CopySnapshot(&ec2.CopySnapshotInput{
+		SourceSnapshotId: snap.SnapshotId,
+	})
+	require.NoError(t, err)
+
+	// Both snapshot IDs should be in KV
+	entry, err := kv.Get(volumeID)
+	require.NoError(t, err)
+	var snapshots []string
+	require.NoError(t, json.Unmarshal(entry.Value(), &snapshots))
+	assert.Contains(t, snapshots, *snap.SnapshotId)
+	assert.Contains(t, snapshots, *copyResult.SnapshotId)
+	assert.Len(t, snapshots, 2)
+}

hive/handlers/ec2/volume/service_impl.go

@@ -34,10 +34,13 @@ type VolumeServiceImpl struct {
 	store      objectstore.ObjectStore
 	bucketName string
 	natsConn   *nats.Conn
+	snapshotKV nats.KeyValue
 }
 
-// NewVolumeServiceImpl creates a new daemon-side volume service
-func NewVolumeServiceImpl(cfg *config.Config, natsConn *nats.Conn) *VolumeServiceImpl {
+// NewVolumeServiceImpl creates a new daemon-side volume service.
+// snapshotKV is optional — when non-nil, DeleteVolume uses O(1) KV lookup
+// instead of scanning all snapshots in S3.
+func NewVolumeServiceImpl(cfg *config.Config, natsConn *nats.Conn, snapshotKV nats.KeyValue) *VolumeServiceImpl {
 	store := objectstore.NewS3ObjectStoreFromConfig(
 		cfg.Predastore.Host,
 		cfg.Predastore.Region,
@@ -50,21 +53,26 @@ func NewVolumeServiceImpl(cfg *config.Config, natsConn *nats.Conn) *VolumeServic
 		store:      store,
 		bucketName: cfg.Predastore.Bucket,
 		natsConn:   natsConn,
+		snapshotKV: snapshotKV,
 	}
 }
 
 // NewVolumeServiceImplWithStore creates a volume service with a custom ObjectStore (for testing)
-func NewVolumeServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn) *VolumeServiceImpl {
+func NewVolumeServiceImplWithStore(cfg *config.Config, store objectstore.ObjectStore, natsConn *nats.Conn, snapshotKV ...nats.KeyValue) *VolumeServiceImpl {
 	bucketName := ""
 	if cfg != nil {
 		bucketName = cfg.Predastore.Bucket
 	}
-	return &VolumeServiceImpl{
+	svc := &VolumeServiceImpl{
 		config:     cfg,
 		store:      store,
 		bucketName: bucketName,
 		natsConn:   natsConn,
 	}
+	if len(snapshotKV) > 0 {
+		svc.snapshotKV = snapshotKV[0]
+	}
+	return svc
 }
 
 // CreateVolume creates a new EBS volume via viperblock and persists its config to S3
@@ -705,9 +713,9 @@ func (s *VolumeServiceImpl) DeleteVolume(input *ec2.DeleteVolumeInput) (*ec2.Del
 		return nil, errors.New(awserrors.ErrorVolumeInUse)
 	}
 
-	// Check if any snapshots reference this volume. Snapshot-backed clones
-	// read chunk files from the source volume's S3 prefix via ReadFrom().
-	// Deleting the source volume would silently break all clones.
+	// Check if any snapshots reference this volume via JetStream KV.
+	// Snapshot-backed clones read chunk files from the source volume's
+	// S3 prefix via ReadFrom(). Deleting the source would break all clones.
 	if err := s.checkVolumeHasNoSnapshots(volumeID); err != nil {
 		return nil, err
 	}
@@ -831,90 +839,40 @@ func (s *VolumeServiceImpl) getSnapshotMetadata(snapshotID string) (*snapshotMet
 	return &meta, nil
 }
 
-// snapshotVolumeRef reads VolumeID from hive metadata.json (volume_id)
-// or viperblock config.json (SourceVolumeName).
-type snapshotVolumeRef struct {
-	VolumeID         string `json:"volume_id"`
-	SourceVolumeName string `json:"SourceVolumeName"`
-}
-
-func (r snapshotVolumeRef) referencesVolume(volumeID string) bool {
-	return r.VolumeID == volumeID || r.SourceVolumeName == volumeID
-}
-
-// checkVolumeHasNoSnapshots scans all snap-* prefixes in S3 and returns an
-// error if any snapshot references the given volume. Checks both hive's
-// metadata.json and viperblock's config.json to handle either format.
+// checkVolumeHasNoSnapshots checks if a volume has dependent snapshots
+// using the JetStream KV index.
 func (s *VolumeServiceImpl) checkVolumeHasNoSnapshots(volumeID string) error {
-	listResult, err := s.store.ListObjects(&s3.ListObjectsInput{
-		Bucket:    aws.String(s.bucketName),
-		Prefix:    aws.String("snap-"),
-		Delimiter: aws.String("/"),
-	})
-	if err != nil {
-		slog.Error("checkVolumeHasNoSnapshots: failed to list snapshots", "volumeId", volumeID, "err", err)
+	if s.snapshotKV == nil {
+		slog.Error("checkVolumeHasNoSnapshots: snapshotKV is nil", "volumeId", volumeID)
 		return errors.New(awserrors.ErrorServerInternal)
 	}
 
-	for _, prefix := range listResult.CommonPrefixes {
-		if prefix.Prefix == nil {
-			continue
-		}
-
-		snapshotID := strings.TrimSuffix(*prefix.Prefix, "/")
-
-		if err := s.snapshotReferencesVolume(snapshotID, volumeID); err != nil {
-			return err
-		}
+	has, err := s.volumeHasSnapshotsKV(volumeID)
+	if err != nil {
+		slog.Error("checkVolumeHasNoSnapshots: KV lookup failed", "volumeId", volumeID, "err", err)
+		return errors.New(awserrors.ErrorServerInternal)
+	}
+	if has {
+		slog.Error("DeleteVolume blocked: volume has snapshots", "volumeId", volumeID)
+		return errors.New(awserrors.ErrorVolumeInUse)
 	}
-
 	return nil
 }
 
-// snapshotReferencesVolume checks if a snapshot references the given volume.
-// Returns an error if the snapshot references the volume or if a non-recoverable
-// S3/IO error occurs. Returns nil if the snapshot does not reference the volume.
-func (s *VolumeServiceImpl) snapshotReferencesVolume(snapshotID, volumeID string) error {
-	// Try metadata.json first (hive format), then config.json (viperblock format)
-	for _, filename := range []string{"metadata.json", "config.json"} {
-		key := snapshotID + "/" + filename
-
-		getResult, err := s.store.GetObject(&s3.GetObjectInput{
-			Bucket: aws.String(s.bucketName),
-			Key:    aws.String(key),
-		})
-		if err != nil {
-			if objectstore.IsNoSuchKeyError(err) {
-				continue // Expected: this format file doesn't exist for this snapshot
-			}
-			slog.Error("checkVolumeHasNoSnapshots: failed to read snapshot metadata",
-				"volumeId", volumeID, "snapshotId", snapshotID, "key", key, "err", err)
-			return errors.New(awserrors.ErrorServerInternal)
-		}
-
-		body, err := io.ReadAll(getResult.Body)
-		if closeErr := getResult.Body.Close(); closeErr != nil {
-			slog.Error("checkVolumeHasNoSnapshots: failed to close response body",
-				"volumeId", volumeID, "snapshotId", snapshotID, "key", key, "err", closeErr)
-		}
-		if err != nil {
-			slog.Error("checkVolumeHasNoSnapshots: failed to read snapshot body",
-				"volumeId", volumeID, "snapshotId", snapshotID, "key", key, "err", err)
-			return errors.New(awserrors.ErrorServerInternal)
-		}
-
-		var ref snapshotVolumeRef
-		if err := json.Unmarshal(body, &ref); err != nil {
-			slog.Warn("checkVolumeHasNoSnapshots: failed to parse snapshot metadata, skipping file",
-				"volumeId", volumeID, "snapshotId", snapshotID, "key", key, "err", err)
-			continue // JSON parse errors for individual files are not fatal -- try next file
-		}
-
-		if ref.referencesVolume(volumeID) {
-			slog.Error("DeleteVolume blocked: volume has snapshots", "volumeId", volumeID, "snapshotId", snapshotID)
-			return fmt.Errorf("%s: volume %s has existing snapshot %s. Delete snapshots first", awserrors.ErrorVolumeInUse, volumeID, snapshotID)
+// volumeHasSnapshotsKV checks the JetStream KV index for snapshot references.
+func (s *VolumeServiceImpl) volumeHasSnapshotsKV(volumeID string) (bool, error) {
+	entry, err := s.snapshotKV.Get(volumeID)
+	if err != nil {
+		if errors.Is(err, nats.ErrKeyNotFound) {
+			return false, nil
 		}
+		return false, err
 	}
 
-	return nil
+	var snapshots []string
+	if err := json.Unmarshal(entry.Value(), &snapshots); err != nil {
+		return false, err
+	}
+
+	return len(snapshots) > 0, nil
 }

hive/handlers/ec2/volume/service_impl_test.go

@@ -4,6 +4,7 @@ import (
 	"encoding/json"
 	"strings"
 	"testing"
+	"time"
 
 	"github.com/aws/aws-sdk-go/aws"
 	"github.com/aws/aws-sdk-go/service/ec2"
@@ -12,6 +13,8 @@ import (
 	"github.com/mulgadc/hive/hive/config"
 	"github.com/mulgadc/hive/hive/objectstore"
 	"github.com/mulgadc/viperblock/viperblock"
+	"github.com/nats-io/nats-server/v2/server"
+	"github.com/nats-io/nats.go"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
 )
@@ -271,98 +274,6 @@ func newTestVolumeServiceWithStore(az string, store *objectstore.MemoryObjectSto
 	return NewVolumeServiceImplWithStore(cfg, store, nil)
 }
 
-func TestDeleteVolume_BlockedBySnapshot(t *testing.T) {
-	store := objectstore.NewMemoryObjectStore()
-	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
-
-	volumeID := "vol-test123"
-
-	// Create volume config in store
-	volumeState := viperblock.VBState{
-		VolumeConfig: viperblock.VolumeConfig{
-			VolumeMetadata: viperblock.VolumeMetadata{
-				VolumeID: volumeID,
-				SizeGiB:  10,
-				State:    "available",
-			},
-		},
-	}
-	data, err := json.Marshal(volumeState)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String(volumeID + "/config.json"),
-		Body:   strings.NewReader(string(data)),
-	})
-	require.NoError(t, err)
-
-	// Create a snapshot referencing this volume
-	snapCfg := snapshotVolumeRef{VolumeID: volumeID}
-	snapData, err := json.Marshal(snapCfg)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String("snap-abc123/config.json"),
-		Body:   strings.NewReader(string(snapData)),
-	})
-	require.NoError(t, err)
-
-	// DeleteVolume should be blocked
-	_, err = svc.DeleteVolume(&ec2.DeleteVolumeInput{
-		VolumeId: aws.String(volumeID),
-	})
-	require.Error(t, err)
-	assert.Contains(t, err.Error(), awserrors.ErrorVolumeInUse)
-	assert.Contains(t, err.Error(), "snap-abc123")
-}
-
-func TestDeleteVolume_AllowedWithoutSnapshots(t *testing.T) {
-	store := objectstore.NewMemoryObjectStore()
-	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
-
-	volumeID := "vol-test456"
-
-	// Create volume config in store
-	volumeState := viperblock.VBState{
-		VolumeConfig: viperblock.VolumeConfig{
-			VolumeMetadata: viperblock.VolumeMetadata{
-				VolumeID: volumeID,
-				SizeGiB:  10,
-				State:    "available",
-			},
-		},
-	}
-	data, err := json.Marshal(volumeState)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String(volumeID + "/config.json"),
-		Body:   strings.NewReader(string(data)),
-	})
-	require.NoError(t, err)
-
-	// Create a snapshot referencing a DIFFERENT volume
-	snapCfg := snapshotVolumeRef{VolumeID: "vol-other"}
-	snapData, err := json.Marshal(snapCfg)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String("snap-xyz789/config.json"),
-		Body:   strings.NewReader(string(snapData)),
-	})
-	require.NoError(t, err)
-
-	// DeleteVolume should succeed (no snapshots reference this volume)
-	_, err = svc.DeleteVolume(&ec2.DeleteVolumeInput{
-		VolumeId: aws.String(volumeID),
-	})
-	require.NoError(t, err)
-}
-
 func TestCreateVolume_FromSnapshot_PassesValidation(t *testing.T) {
 	store := objectstore.NewMemoryObjectStore()
 	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
@@ -502,100 +413,6 @@ func TestCreateVolume_FromSnapshot_SizeEqualToSnapshot(t *testing.T) {
 	}
 }
 
-func TestDeleteVolume_BlockedBySourceVolumeName(t *testing.T) {
-	store := objectstore.NewMemoryObjectStore()
-	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
-
-	volumeID := "vol-srcname"
-
-	// Create volume config
-	volumeState := viperblock.VBState{
-		VolumeConfig: viperblock.VolumeConfig{
-			VolumeMetadata: viperblock.VolumeMetadata{
-				VolumeID: volumeID,
-				SizeGiB:  10,
-				State:    "available",
-			},
-		},
-	}
-	data, err := json.Marshal(volumeState)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String(volumeID + "/config.json"),
-		Body:   strings.NewReader(string(data)),
-	})
-	require.NoError(t, err)
-
-	// Snapshot references volume via SourceVolumeName (viperblock config.json format)
-	snapCfg := snapshotVolumeRef{SourceVolumeName: volumeID}
-	snapData, err := json.Marshal(snapCfg)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String("snap-vb001/config.json"),
-		Body:   strings.NewReader(string(snapData)),
-	})
-	require.NoError(t, err)
-
-	// DeleteVolume should be blocked via SourceVolumeName match
-	_, err = svc.DeleteVolume(&ec2.DeleteVolumeInput{
-		VolumeId: aws.String(volumeID),
-	})
-	require.Error(t, err)
-	assert.Contains(t, err.Error(), awserrors.ErrorVolumeInUse)
-	assert.Contains(t, err.Error(), "snap-vb001")
-}
-
-func TestDeleteVolume_BlockedByMetadataJson(t *testing.T) {
-	store := objectstore.NewMemoryObjectStore()
-	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
-
-	volumeID := "vol-metablock"
-
-	// Create volume config
-	volumeState := viperblock.VBState{
-		VolumeConfig: viperblock.VolumeConfig{
-			VolumeMetadata: viperblock.VolumeMetadata{
-				VolumeID: volumeID,
-				SizeGiB:  10,
-				State:    "available",
-			},
-		},
-	}
-	data, err := json.Marshal(volumeState)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String(volumeID + "/config.json"),
-		Body:   strings.NewReader(string(data)),
-	})
-	require.NoError(t, err)
-
-	// Snapshot references volume via metadata.json (hive format, volume_id field)
-	snapCfg := snapshotVolumeRef{VolumeID: volumeID}
-	snapData, err := json.Marshal(snapCfg)
-	require.NoError(t, err)
-
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String("snap-meta001/metadata.json"),
-		Body:   strings.NewReader(string(snapData)),
-	})
-	require.NoError(t, err)
-
-	// DeleteVolume should be blocked via metadata.json match
-	_, err = svc.DeleteVolume(&ec2.DeleteVolumeInput{
-		VolumeId: aws.String(volumeID),
-	})
-	require.Error(t, err)
-	assert.Contains(t, err.Error(), awserrors.ErrorVolumeInUse)
-	assert.Contains(t, err.Error(), "snap-meta001")
-}
-
 func TestCreateVolume_FromSnapshot_CorruptMetadata(t *testing.T) {
 	store := objectstore.NewMemoryObjectStore()
 	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
@@ -617,13 +434,39 @@ func TestCreateVolume_FromSnapshot_CorruptMetadata(t *testing.T) {
 	require.Error(t, err)
 }
 
-func TestDeleteVolume_CorruptSnapshotMetadata_SkipsFile(t *testing.T) {
-	store := objectstore.NewMemoryObjectStore()
-	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
+// setupTestVolumeKV creates a NATS JetStream test server and returns a KV bucket.
+func setupTestVolumeKV(t *testing.T) nats.KeyValue {
+	t.Helper()
+	opts := &server.Options{
+		Host:      "127.0.0.1",
+		Port:      -1,
+		JetStream: true,
+		StoreDir:  t.TempDir(),
+		NoLog:     true,
+		NoSigs:    true,
+	}
+	ns, err := server.NewServer(opts)
+	require.NoError(t, err)
+	go ns.Start()
+	require.True(t, ns.ReadyForConnections(5*time.Second))
+	t.Cleanup(func() { ns.Shutdown() })
 
-	volumeID := "vol-corruptsnap"
+	nc, err := nats.Connect(ns.ClientURL())
+	require.NoError(t, err)
+	t.Cleanup(func() { nc.Close() })
 
-	// Create volume config
+	js, err := nc.JetStream()
+	require.NoError(t, err)
+
+	kv, err := js.CreateKeyValue(&nats.KeyValueConfig{
+		Bucket: "hive-volume-snapshots",
+	})
+	require.NoError(t, err)
+	return kv
+}
+
+func createVolumeInStore(t *testing.T, store *objectstore.MemoryObjectStore, volumeID string) {
+	t.Helper()
 	volumeState := viperblock.VBState{
 		VolumeConfig: viperblock.VolumeConfig{
 			VolumeMetadata: viperblock.VolumeMetadata{
@@ -642,19 +485,59 @@ func TestDeleteVolume_CorruptSnapshotMetadata_SkipsFile(t *testing.T) {
 		Body:   strings.NewReader(string(data)),
 	})
 	require.NoError(t, err)
+}
 
-	// Put corrupt metadata.json (invalid JSON) -- should be skipped
-	_, err = store.PutObject(&s3.PutObjectInput{
-		Bucket: aws.String("test-bucket"),
-		Key:    aws.String("snap-corrupt/metadata.json"),
-		Body:   strings.NewReader("not json{{{"),
-	})
+func TestDeleteVolume_BlockedByKV(t *testing.T) {
+	kv := setupTestVolumeKV(t)
+	store := objectstore.NewMemoryObjectStore()
+	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
+	svc.snapshotKV = kv
+
+	volumeID := "vol-kvblocked"
+	createVolumeInStore(t, store, volumeID)
+
+	// Put a snapshot ref in KV
+	data, err := json.Marshal([]string{"snap-001"})
+	require.NoError(t, err)
+	_, err = kv.Put(volumeID, data)
 	require.NoError(t, err)
 
-	// No valid config.json either, so no snapshot blocks this volume.
-	// DeleteVolume should succeed since corrupt metadata is skipped.
+	// DeleteVolume should be blocked
 	_, err = svc.DeleteVolume(&ec2.DeleteVolumeInput{
 		VolumeId: aws.String(volumeID),
 	})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), awserrors.ErrorVolumeInUse)
+}
+
+func TestDeleteVolume_AllowedByKV(t *testing.T) {
+	kv := setupTestVolumeKV(t)
+	store := objectstore.NewMemoryObjectStore()
+	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
+	svc.snapshotKV = kv
+
+	volumeID := "vol-kvallowed"
+	createVolumeInStore(t, store, volumeID)
+
+	// No KV entry → delete allowed
+	_, err := svc.DeleteVolume(&ec2.DeleteVolumeInput{
+		VolumeId: aws.String(volumeID),
+	})
 	require.NoError(t, err)
 }
+
+func TestDeleteVolume_ErrorWhenKVNil(t *testing.T) {
+	store := objectstore.NewMemoryObjectStore()
+	svc := newTestVolumeServiceWithStore("ap-southeast-2a", store)
+	// snapshotKV is nil by default
+
+	volumeID := "vol-nokvtest"
+	createVolumeInStore(t, store, volumeID)
+
+	// Should fail because snapshotKV is nil
+	_, err := svc.DeleteVolume(&ec2.DeleteVolumeInput{
+		VolumeId: aws.String(volumeID),
+	})
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), awserrors.ErrorServerInternal)
+}