From 6f30300639b0f3be443278c9d9e7723a5bead56a Mon Sep 17 00:00:00 2001
From: David Son <davbson@amazon.com>
Date: Sat, 13 Jan 2024 00:59:43 +0000
Subject: [PATCH] Keep directories when SIGINT sent to daemon

Signed-off-by: David Son <davbson@amazon.com>
---
 snapshot/snapshot.go | 37 ++++++++++++++++++++++++++++---------
 1 file changed, 28 insertions(+), 9 deletions(-)

diff --git a/snapshot/snapshot.go b/snapshot/snapshot.go
index 1a29f1ce8..b090cc202 100644
--- a/snapshot/snapshot.go
+++ b/snapshot/snapshot.go
@@ -471,13 +471,8 @@ func (o *snapshotter) getCleanupDirectories(ctx context.Context, t storage.Trans
 }
 
 func (o *snapshotter) cleanupSnapshotDirectory(ctx context.Context, dir string) error {
-
-	// On a remote snapshot, the layer is mounted on the "fs" directory.
-	// We use Filesystem's Unmount API so that it can do necessary finalization
-	// before/after the unmount.
-	mp := filepath.Join(dir, "fs")
-	if err := o.fs.Unmount(ctx, mp); err != nil {
-		log.G(ctx).WithError(err).WithField("dir", mp).Debug("failed to unmount")
+	if err := o.unmountSnapshotDirectory(ctx, dir); err != nil {
+		log.G(ctx).WithError(err).WithField("dir", dir).Debug("failed to unmount")
 	}
 	if err := os.RemoveAll(dir); err != nil {
 		return fmt.Errorf("failed to remove directory %q: %w", dir, err)
@@ -485,6 +480,14 @@ func (o *snapshotter) cleanupSnapshotDirectory(ctx context.Context, dir string)
 	return nil
 }
 
+func (o *snapshotter) unmountSnapshotDirectory(ctx context.Context, dir string) error {
+	// On a remote snapshot, the layer is mounted on the "fs" directory.
+	// We use Filesystem's Unmount API so that it can do necessary finalization
+	// before/after the unmount.
+	mp := filepath.Join(dir, "fs")
+	return o.fs.Unmount(ctx, mp)
+}
+
 func (o *snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, key, parent string, opts []snapshots.Opt) (_ storage.Snapshot, err error) {
 	ctx, t, err := o.ms.TransactionContext(ctx, true)
 	if err != nil {
@@ -656,12 +659,28 @@ func (o *snapshotter) Close() error {
 	// unmount all mounts including Committed
 	const cleanupCommitted = true
 	ctx := context.Background()
-	if err := o.cleanup(ctx, cleanupCommitted); err != nil {
-		log.G(ctx).WithError(err).Warn("failed to cleanup")
+	if err := o.unmountAllSnapshots(ctx, cleanupCommitted); err != nil {
+		log.G(ctx).WithError(err).Warn("failed to unmount snapshots on close")
 	}
 	return o.ms.Close()
 }
 
+func (o *snapshotter) unmountAllSnapshots(ctx context.Context, cleanupCommitted bool) error {
+	cleanup, err := o.cleanupDirectories(ctx, cleanupCommitted)
+	if err != nil {
+		return err
+	}
+
+	log.G(ctx).Debugf("unmount: dirs=%v", cleanup)
+	for _, dir := range cleanup {
+		if err := o.unmountSnapshotDirectory(ctx, dir); err != nil {
+			log.G(ctx).WithError(err).WithField("path", dir).Warn("failed to unmount directory")
+		}
+	}
+
+	return nil
+}
+
 // prepareRemoteSnapshot tries to prepare the snapshot as a remote snapshot
 // using filesystems registered in this snapshotter.
 func (o *snapshotter) prepareRemoteSnapshot(ctx context.Context, key string, labels map[string]string) error {