Keep directories when SIGINT sent to daemon

Signed-off-by: David Son <davbson@amazon.com>

snapshot/snapshot.go

@@ -467,7 +467,7 @@ func (o *snapshotter) Remove(ctx context.Context, key string) (err error) {
 		defer func() {
 			if err == nil {
 				for _, dir := range removals {
-					if err := o.cleanupSnapshotDirectory(ctx, dir); err != nil {
+					if err := o.cleanupSnapshotDirectory(ctx, dir, true); err != nil {
 						log.G(ctx).WithError(err).WithField("path", dir).Warn("failed to remove directory")
 					}
 				}
@@ -492,18 +492,18 @@ func (o *snapshotter) Walk(ctx context.Context, fn snapshots.WalkFunc, fs ...str
 // Cleanup cleans up disk resources from removed or abandoned snapshots
 func (o *snapshotter) Cleanup(ctx context.Context) error {
 	const cleanupCommitted = false
-	return o.cleanup(ctx, cleanupCommitted)
+	return o.cleanup(ctx, cleanupCommitted, true)
 }
 
-func (o *snapshotter) cleanup(ctx context.Context, cleanupCommitted bool) error {
+func (o *snapshotter) cleanup(ctx context.Context, cleanupCommitted bool, removeDir bool) error {
 	cleanup, err := o.cleanupDirectories(ctx, cleanupCommitted)
 	if err != nil {
 		return err
 	}
 
 	log.G(ctx).Debugf("cleanup: dirs=%v", cleanup)
 	for _, dir := range cleanup {
-		if err := o.cleanupSnapshotDirectory(ctx, dir); err != nil {
+		if err := o.cleanupSnapshotDirectory(ctx, dir, removeDir); err != nil {
 			log.G(ctx).WithError(err).WithField("path", dir).Warn("failed to remove directory")
 		}
 	}
@@ -555,18 +555,24 @@ func (o *snapshotter) getCleanupDirectories(ctx context.Context, t storage.Trans
 	return cleanup, nil
 }
 
-func (o *snapshotter) cleanupSnapshotDirectory(ctx context.Context, dir string) error {
+func (o *snapshotter) cleanupSnapshotDirectory(ctx context.Context, dir string, removeDir bool) error {
+	o.unmountSnapshotDirectory(ctx, dir)
+	if removeDir {
+		if err := os.RemoveAll(dir); err != nil {
+			return fmt.Errorf("failed to remove directory %q: %w", dir, err)
+		}
+	}
+	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")
 	if err := o.fs.Unmount(ctx, mp); err != nil {
 		log.G(ctx).WithError(err).WithField("dir", mp).Debug("failed to unmount")
 	}
-	if err := os.RemoveAll(dir); err != nil {
-		return fmt.Errorf("failed to remove directory %q: %w", dir, err)
-	}
 	return nil
 }
 
@@ -580,12 +586,12 @@ func (o *snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, k
 	defer func() {
 		if err != nil {
 			if td != "" {
-				if err1 := o.cleanupSnapshotDirectory(ctx, td); err1 != nil {
+				if err1 := o.cleanupSnapshotDirectory(ctx, td, true); err1 != nil {
 					log.G(ctx).WithError(err1).Warn("failed to cleanup temp snapshot directory")
 				}
 			}
 			if path != "" {
-				if err1 := o.cleanupSnapshotDirectory(ctx, path); err1 != nil {
+				if err1 := o.cleanupSnapshotDirectory(ctx, path, true); err1 != nil {
 					log.G(ctx).WithError(err1).WithField("path", path).Error("failed to reclaim snapshot directory, directory may need removal")
 					err = fmt.Errorf("failed to remove path: %v: %w", err1, err)
 				}
@@ -743,7 +749,7 @@ func (o *snapshotter) Close() error {
 	// unmount all mounts including Committed
 	const cleanupCommitted = true
 	ctx := context.Background()
-	if err := o.cleanup(ctx, cleanupCommitted); err != nil {
+	if err := o.cleanup(ctx, cleanupCommitted, false); err != nil {
 		log.G(ctx).WithError(err).Warn("failed to cleanup")
 	}
 	return o.ms.Close()