From 7529c33839914f3811b3ac3cff87fcdf90c4971d Mon Sep 17 00:00:00 2001
From: David Son <davbson@amazon.com>
Date: Fri, 20 Oct 2023 22:47:49 +0000
Subject: [PATCH] Keep directories when SIGINT sent to daemon

Signed-off-by: David Son <davbson@amazon.com>
---
 integration/run_test.go | 30 ++++++++++++++++++++++++++++++
 snapshot/snapshot.go    | 32 +++++++++++++++++++++++++++-----
 2 files changed, 57 insertions(+), 5 deletions(-)

diff --git a/integration/run_test.go b/integration/run_test.go
index d46092fc8..ca2324655 100644
--- a/integration/run_test.go
+++ b/integration/run_test.go
@@ -424,6 +424,36 @@ func TestRootFolderPermission(t *testing.T) {
 	}
 }
 
+func TestRestartAfterSigint(t *testing.T) {
+	const containerImage = alpineImage
+
+	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") // pkill can only take up to 15 chars
+
+	timeout := []string{"timeout", "--preserve-status", "5s"} // arbitrarily chosen value
+	cmd := shell.C("/usr/local/bin/soci-snapshotter-grpc", "--log-level", sociLogLevel,
+		"--address", "/run/soci-snapshotter-grpc/soci-snapshotter-grpc.sock")
+	cmd = addConfig(t, sh, getSnapshotterConfigToml(t, false, tcpMetricsConfig), cmd...)
+	cmd = append(timeout, cmd...)
+
+	// exit status 0 — command timed out, so server is running — PASS
+	// exit status 1 — server crashed before time limit — FAIL
+	if _, err := sh.OLog(cmd...); err != nil {
+		if err.Error() == "exit status 1" {
+			t.Logf("error starting snapshotter daemon")
+			t.Fail()
+		}
+	}
+}
+
 func TestRunInContentStore(t *testing.T) {
 	imageName := helloImage
 	sh, done := newSnapshotterBaseShell(t)
diff --git a/snapshot/snapshot.go b/snapshot/snapshot.go
index b80d8d901..5a0298fa3 100644
--- a/snapshot/snapshot.go
+++ b/snapshot/snapshot.go
@@ -556,7 +556,16 @@ 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.
@@ -564,9 +573,6 @@ func (o *snapshotter) cleanupSnapshotDirectory(ctx context.Context, dir string)
 	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)