Keep directories when SIGINT sent to daemon

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

snapshot/snapshot.go

@@ -134,6 +134,11 @@ func AllowInvalidMountsOnRestart(config *SnapshotterConfig) error {
 	return nil
 }
 
+type CleanupConfig struct {
+	CleanupCommitted bool
+	RemoveDir        bool
+}
+
 type snapshotter struct {
 	root        string
 	ms          *storage.MetaStore
@@ -466,8 +471,12 @@ func (o *snapshotter) Remove(ctx context.Context, key string) (err error) {
 		// key no longer available.
 		defer func() {
 			if err == nil {
+				cleanupCfg := CleanupConfig{
+					CleanupCommitted: true,
+					RemoveDir:        true,
+				}
 				for _, dir := range removals {
-					if err := o.cleanupSnapshotDirectory(ctx, dir); err != nil {
+					if err := o.cleanupSnapshotDirectory(ctx, dir, &cleanupCfg); err != nil {
 						log.G(ctx).WithError(err).WithField("path", dir).Warn("failed to remove directory")
 					}
 				}
@@ -491,19 +500,22 @@ 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)
+	cleanupCfg := CleanupConfig{
+		CleanupCommitted: false,
+		RemoveDir:        true,
+	}
+	return o.cleanup(ctx, &cleanupCfg)
 }
 
-func (o *snapshotter) cleanup(ctx context.Context, cleanupCommitted bool) error {
-	cleanup, err := o.cleanupDirectories(ctx, cleanupCommitted)
+func (o *snapshotter) cleanup(ctx context.Context, cleanupCfg *CleanupConfig) error {
+	cleanup, err := o.cleanupDirectories(ctx, cleanupCfg.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, cleanupCfg); err != nil {
 			log.G(ctx).WithError(err).WithField("path", dir).Warn("failed to remove directory")
 		}
 	}
@@ -555,18 +567,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, cleanupCfg *CleanupConfig) error {
+	o.unmountSnapshotDirectory(ctx, dir)
+	if cleanupCfg.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
 }
 
@@ -579,13 +597,17 @@ func (o *snapshotter) createSnapshot(ctx context.Context, kind snapshots.Kind, k
 	var td, path string
 	defer func() {
 		if err != nil {
+			cleanupCfg := CleanupConfig{
+				CleanupCommitted: true,
+				RemoveDir:        true,
+			}
 			if td != "" {
-				if err1 := o.cleanupSnapshotDirectory(ctx, td); err1 != nil {
+				if err1 := o.cleanupSnapshotDirectory(ctx, td, &cleanupCfg); 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, &cleanupCfg); 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)
 				}
@@ -741,9 +763,12 @@ func (o *snapshotter) workPath(id string) string {
 // Close closes the snapshotter
 func (o *snapshotter) Close() error {
 	// unmount all mounts including Committed
-	const cleanupCommitted = true
+	cleanupCfg := CleanupConfig{
+		CleanupCommitted: true,
+		RemoveDir:        false,
+	}
 	ctx := context.Background()
-	if err := o.cleanup(ctx, cleanupCommitted); err != nil {
+	if err := o.cleanup(ctx, &cleanupCfg); err != nil {
 		log.G(ctx).WithError(err).Warn("failed to cleanup")
 	}
 	return o.ms.Close()