Keep directories when SIGINT sent to daemon

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

integration/run_test.go

@@ -424,6 +424,51 @@ func TestRootFolderPermission(t *testing.T) {
 	}
 }
 
+func TestRestartAfterSigint(t *testing.T) {
+	const containerImage = alpineImage
+	const killTimeout = 5
+	const startTimeout = "5s"
+
+	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", "-SIGINT", "soci-snapshotte") // pkill can only take up to 15 chars
+
+	var buffer []byte
+	timedOut := true
+	for i := 0; i < killTimeout*4; i++ {
+		buffer = sh.O("ps")
+		if !strings.Contains(string(buffer), "soci-snapshotte") {
+			timedOut = false
+			break
+		}
+		sh.X("sleep", "0.25")
+	}
+
+	if timedOut {
+		t.Fatalf("failed to kill snapshotter daemon")
+	}
+
+	timeoutCmd := []string{"timeout", "--preserve-status", startTimeout}
+	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(timeoutCmd, 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 0" {
+			t.Fatalf("error starting snapshotter daemon: %v", err)
+		}
+	}
+}
+
 func TestRunInContentStore(t *testing.T) {
 	imageName := helloImage
 	sh, done := newSnapshotterBaseShell(t)

snapshot/snapshot.go

@@ -556,20 +556,23 @@ 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 {
+		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")
+	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 {
@@ -743,12 +746,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)