fix: clean up partial containers when interrupted (#2090)

Wire signal.NotifyContext in main so SIGINT/SIGTERM cancel the context
instead of killing the process.

ephemeral and assemble now defer rm on a detached cleanup ctx, so the
in-flight container is removed even when the user hits Ctrl+C.

Signed-off-by: Luca Di Maio <luca.dimaio1@gmail.com>

cmd/distrobox/main.go

@@ -4,19 +4,32 @@ import (
 	"context"
 	"log"
 	"os"
+	"os/signal"
+	"syscall"
 
 	"github.com/89luca89/distrobox/internal/cli"
 	"github.com/89luca89/distrobox/pkg/config"
 )
 
 func main() {
-	cfg, err := config.LoadValues()
-	if err != nil {
-		log.Fatal(err)
-	}
-
-	cmd := cli.NewRootCommand(cfg)
-	if err := cmd.Run(context.Background(), os.Args); err != nil {
+	if err := run(); err != nil {
 		log.Fatal(err)
 	}
 }
+
+func run() error {
+	cfg, err := config.LoadValues()
+	if err != nil {
+		//nolint:wrapcheck // main reports errors as-is
+		return err
+	}
+
+	// SIGINT register
+	ctx, stop := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
+	defer stop()
+
+	cmd := cli.NewRootCommand(cfg)
+
+	//nolint:wrapcheck // main reports errors as-is
+	return cmd.Run(ctx, os.Args)
+}

pkg/commands/assemble.go

@@ -6,6 +6,7 @@ import (
 	"regexp"
 	"slices"
 	"strings"
+	"time"
 
 	"github.com/89luca89/distrobox/pkg/config"
 	"github.com/89luca89/distrobox/pkg/containermanager"
@@ -13,6 +14,8 @@ import (
 	"github.com/89luca89/distrobox/pkg/ui"
 )
 
+const assembleCleanupTimeout = 30 * time.Second
+
 type AssembleOptions struct {
 	Items []manifest.Item
 	// Boxname is the name of the box to assemble
@@ -36,6 +39,7 @@ type AssembleCommand struct {
 	rmCmdRoot     *RmCommand
 	enterCmdRoot  *EnterCommand
 	progress      *ui.Progress
+	printer       *ui.Printer
 }
 
 func NewAssembleCommand(
@@ -55,6 +59,7 @@ func NewAssembleCommand(
 		rmCmdRoot:     NewRmCommand(cfg, cmRoot, prompter),
 		enterCmdRoot:  NewEnterCommand(cfg, cmRoot, progress, printer),
 		progress:      progress,
+		printer:       printer,
 	}
 }
 
@@ -162,6 +167,26 @@ func (ac *AssembleCommand) createItem(ctx context.Context, item manifest.Item, d
 		return err
 	}
 
+	success := false
+	defer func() {
+		if success || dryRun {
+			return
+		}
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), assembleCleanupTimeout)
+		defer cancel()
+		rmCmd := ac.rmCmd
+		if item.Root {
+			rmCmd = ac.rmCmdRoot
+		}
+		if _, rmErr := rmCmd.Execute(cleanupCtx, RmOptions{
+			NoTTY:          true,
+			Force:          true,
+			ContainerNames: []string{item.Name},
+		}); rmErr != nil {
+			ac.printer.PrintWarningln("warning: %s: %s", item.Name, rmErr)
+		}
+	}()
+
 	if !dryRun {
 		err = ac.setupBox(ctx, item)
 		if err != nil {
@@ -171,6 +196,7 @@ func (ac *AssembleCommand) createItem(ctx context.Context, item manifest.Item, d
 	}
 
 	ac.progress.Done()
+	success = true
 	return nil
 }
 

pkg/commands/ephemeral.go

@@ -4,12 +4,15 @@ import (
 	"context"
 	"fmt"
 	"math/rand/v2"
+	"time"
 
 	"github.com/89luca89/distrobox/pkg/config"
 	"github.com/89luca89/distrobox/pkg/containermanager"
 	"github.com/89luca89/distrobox/pkg/ui"
 )
 
+const ephemeralCleanupTimeout = 30 * time.Second
+
 type EphemeralOptions struct {
 	CreateOptions
 
@@ -22,6 +25,7 @@ type EphemeralCommand struct {
 	createCmd        *CreateCommand
 	enterCmd         *EnterCommand
 	rmCmd            *RmCommand
+	printer          *ui.Printer
 }
 
 func NewEphemeralCommand(
@@ -37,6 +41,7 @@ func NewEphemeralCommand(
 		createCmd:        NewCreateCommand(cfg, cm, progress, prompter),
 		enterCmd:         NewEnterCommand(cfg, cm, progress, printer),
 		rmCmd:            NewRmCommand(cfg, cm, prompter),
+		printer:          printer,
 	}
 }
 
@@ -53,29 +58,31 @@ func (c *EphemeralCommand) Execute(ctx context.Context, opts EphemeralOptions) e
 	createOpts.GenerateEntry = false
 	createOpts.DryRun = opts.DryRun
 	createOpts.NonInteractive = true
-	if _, err := c.createCmd.Execute(ctx, createOpts); err != nil {
-		return fmt.Errorf("failed to create ephemeral container: %w", err)
+	if _, createErr := c.createCmd.Execute(ctx, createOpts); createErr != nil {
+		return fmt.Errorf("ephemeral: %w", createErr)
 	}
 
+	defer func() {
+		cleanupCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), ephemeralCleanupTimeout)
+		defer cancel()
+		rmOpts := RmOptions{
+			ContainerNames: []string{name},
+			Force:          true,
+			NoTTY:          true,
+		}
+		if _, rmErr := c.rmCmd.Execute(cleanupCtx, rmOpts); rmErr != nil {
+			c.printer.PrintWarningln("warning: %s: %s", name, rmErr)
+		}
+	}()
+
 	// enter into it
 	enterOpts := EnterOptions{
 		ContainerName: name,
 		DryRun:        opts.DryRun,
 		// TODO: handle enter command
 	}
-	if _, err := c.enterCmd.Execute(ctx, enterOpts); err != nil {
-		return fmt.Errorf("failed to enter ephemeral container: %w", err)
-	}
-
-	// remove it
-	rmOpts := RmOptions{
-		ContainerNames: []string{name},
-		Force:          true,
-		NoTTY:          true,
-		All:            false,
-	}
-	if _, err := c.rmCmd.Execute(ctx, rmOpts); err != nil {
-		return fmt.Errorf("failed to remove ephemeral container: %w", err)
+	if _, enterErr := c.enterCmd.Execute(ctx, enterOpts); enterErr != nil {
+		return fmt.Errorf("ephemeral: %w", enterErr)
 	}
 
 	return nil

pkg/containermanager/providers/docker.go

@@ -842,6 +842,10 @@ func (d *Docker) waitForSetup(
 	printer *ui.Printer,
 ) error {
 	for {
+		if err := ctx.Err(); err != nil {
+			return fmt.Errorf("docker: %w", err)
+		}
+
 		// Check container is still running
 		inspectResult, err := d.InspectContainer(ctx, containerName)
 		if err != nil {

pkg/containermanager/providers/podman.go

@@ -923,6 +923,10 @@ func (p *Podman) waitForSetup(
 	printer *ui.Printer,
 ) error {
 	for {
+		if err := ctx.Err(); err != nil {
+			return fmt.Errorf("podman: %w", err)
+		}
+
 		// Check container is still running
 		inspectResult, err := p.InspectContainer(ctx, containerName)
 		if err != nil {