Update metadata snapshotter to lease on exists

dmcgowan · dmcgowan · commit 9f6c61ab927b · 2024-05-09T14:21:57.000-07:00
Currently the metadata snapshotter is not consistently adding keys to a lease when already exists is returned. When a lease is provided, any already exists errors should add the relevant key to the lease. It is not expected that clients must explicitly lease a key after calling Prepare/Commit. Signed-off-by: Derek McGowan <derek@mcg.dev> (cherry picked from commit c7fb8a9) Signed-off-by: Derek McGowan <derek@mcg.dev>
diff --git a/metadata/snapshot.go b/metadata/snapshot.go
@@ -323,6 +323,7 @@ func (s *snapshotter) createSnapshot(ctx context.Context, key, parent string, re
 		bopts   = []snapshots.Opt{
 			snapshots.WithLabels(snapshots.FilterInheritedLabels(base.Labels)),
 		}
+		rerr error
 	)
 
 	if err := update(ctx, s.db, func(tx *bolt.Tx) error {
@@ -334,12 +335,20 @@ func (s *snapshotter) createSnapshot(ctx context.Context, key, parent string, re
 		// Check if target exists, if so, return already exists
 		if target != "" {
 			if tbkt := bkt.Bucket([]byte(target)); tbkt != nil {
-				return fmt.Errorf("target snapshot %q: %w", target, errdefs.ErrAlreadyExists)
+				rerr = fmt.Errorf("target snapshot %q: %w", target, errdefs.ErrAlreadyExists)
+				if err := addSnapshotLease(ctx, tx, s.name, target); err != nil {
+					return err
+				}
+				return nil
 			}
 		}
 
 		if bbkt := bkt.Bucket([]byte(key)); bbkt != nil {
-			return fmt.Errorf("snapshot %q: %w", key, errdefs.ErrAlreadyExists)
+			rerr = fmt.Errorf("snapshot %q: %w", key, errdefs.ErrAlreadyExists)
+			if err := addSnapshotLease(ctx, tx, s.name, key); err != nil {
+				return err
+			}
+			return nil
 		}
 
 		if parent != "" {
@@ -360,11 +369,14 @@ func (s *snapshotter) createSnapshot(ctx context.Context, key, parent string, re
 	}); err != nil {
 		return nil, err
 	}
+	// Already exists and lease successfully added in transaction
+	if rerr != nil {
+		return nil, rerr
+	}
 
 	var (
 		m       []mount.Mount
 		created string
-		rerr    error
 	)
 	if readonly {
 		m, err = s.Snapshotter.View(ctx, bkey, bparent, bopts...)
@@ -527,24 +539,28 @@ func (s *snapshotter) Commit(ctx context.Context, name, key string, opts ...snap
 		return err
 	}
 
-	var bname string
+	var (
+		bname string
+		rerr  error
+	)
 	if err := update(ctx, s.db, func(tx *bolt.Tx) error {
 		bkt := getSnapshotterBucket(tx, ns, s.name)
 		if bkt == nil {
 			return fmt.Errorf("can not find snapshotter %q: %w",
 				s.name, errdefs.ErrNotFound)
 		}
 
+		if err := addSnapshotLease(ctx, tx, s.name, name); err != nil {
+			return err
+		}
 		bbkt, err := bkt.CreateBucket([]byte(name))
 		if err != nil {
 			if err == bolt.ErrBucketExists {
-				err = fmt.Errorf("snapshot %q: %w", name, errdefs.ErrAlreadyExists)
+				rerr = fmt.Errorf("snapshot %q: %w", name, errdefs.ErrAlreadyExists)
+				return nil
 			}
 			return err
 		}
-		if err := addSnapshotLease(ctx, tx, s.name, name); err != nil {
-			return err
-		}
 
 		obkt := bkt.Bucket([]byte(key))
 		if obkt == nil {
@@ -634,7 +650,7 @@ func (s *snapshotter) Commit(ctx context.Context, name, key string, opts ...snap
 		return err
 	}
 
-	if s.db.dbopts.publisher != nil {
+	if rerr == nil && s.db.dbopts.publisher != nil {
 		if err := s.db.dbopts.publisher.Publish(ctx, "/snapshot/commit", &eventstypes.SnapshotCommit{
 			Key:         key,
 			Name:        name,
@@ -644,7 +660,7 @@ func (s *snapshotter) Commit(ctx context.Context, name, key string, opts ...snap
 		}
 	}
 
-	return nil
+	return rerr
 
 }
 
