diff --git a/hive/daemon/daemon.go b/hive/daemon/daemon.go index 441ac7aee..11c2945fb 100644 --- a/hive/daemon/daemon.go +++ b/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 { diff --git a/hive/daemon/daemon_test.go b/hive/daemon/daemon_test.go index 9bee86b6b..62c924f06 100644 --- a/hive/daemon/daemon_test.go +++ b/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 { diff --git a/hive/handlers/ec2/snapshot/service_impl.go b/hive/handlers/ec2/snapshot/service_impl.go index f067d3ea8..1f6312274 100644 --- a/hive/handlers/ec2/snapshot/service_impl.go +++ b/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 +} diff --git a/hive/handlers/ec2/snapshot/service_impl_test.go b/hive/handlers/ec2/snapshot/service_impl_test.go index 81b52fff6..07e00426f 100644 --- a/hive/handlers/ec2/snapshot/service_impl_test.go +++ b/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) +} diff --git a/hive/handlers/ec2/volume/service_impl.go b/hive/handlers/ec2/volume/service_impl.go index 6704dd5d2..47b09e001 100644 --- a/hive/handlers/ec2/volume/service_impl.go +++ b/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 } diff --git a/hive/handlers/ec2/volume/service_impl_test.go b/hive/handlers/ec2/volume/service_impl_test.go index 14a6833a8..09d769c03 100644 --- a/hive/handlers/ec2/volume/service_impl_test.go +++ b/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) +}