Keep directories when SIGINT sent to daemon

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

integration/run_test.go

@@ -424,6 +424,25 @@ func TestRootFolderPermission(t *testing.T) {
 	}
 }
 
+func TestRestartAfterSigint(t *testing.T) {
+	const containerImage = rabbitmqImage
+
+	regConfig := newRegistryConfig()
+	sh, done := newShellWithRegistry(t, regConfig)
+	defer done()
+
+	rebootContainerd(t, sh, getContainerdConfigToml(t, false), getSnapshotterConfigToml(t, false, tcpMetricsConfig))
+
+	copyImage(sh, dockerhub(containerImage), regConfig.mirror(containerImage))
+	indexDigest := buildIndex(sh, regConfig.mirror(containerImage), withMinLayerSize(0), withSpanSize(100*1024))
+	sh.X("soci", "image", "rpull", "--user", regConfig.creds(), "--soci-index-digest", indexDigest, regConfig.mirror(containerImage).ref)
+
+	sh.X("pkill", "-2", "soci-snapshotte")
+	sh.Gox("soci-snapshotter-grpc") // If snapshotter starts, test passes
+	sh.X("sleep", "5")              // Ensure the snapshotter is up and running
+	sh.X("pkill", "-2", "soci-snapshotte")
+}
+
 func TestRunInContentStore(t *testing.T) {
 	imageName := helloImage
 	sh, done := newSnapshotterBaseShell(t)

snapshot/snapshot.go

@@ -556,17 +556,23 @@ func (o *snapshotter) getCleanupDirectories(ctx context.Context, t storage.Trans
 }
 
 func (o *snapshotter) cleanupSnapshotDirectory(ctx context.Context, dir string) error {
+	if err := o.unmountSnapshotDirectory(ctx, dir); err != nil {
+		return fmt.Errorf("failed to unmount directory %q: %w", dir, err)
+	}
+	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
 }
 
@@ -743,12 +749,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
+}
+
 // prepareLocalSnapshot tries to prepare the snapshot as a local snapshot.
 func (o *snapshotter) prepareLocalSnapshot(ctx context.Context, key string, labels map[string]string, mounts []mount.Mount) error {
 	ctx, t, err := o.ms.TransactionContext(ctx, false)