From 75ddc5b811ae620dea857d3c2161aee158cc8f07 Mon Sep 17 00:00:00 2001
From: Tonis Tiigi <tonistiigi@gmail.com>
Date: Sun, 7 Aug 2022 22:16:14 -0700
Subject: [PATCH] build: fix issues with leaving invoke containers running

Signed-off-by: Tonis Tiigi <tonistiigi@gmail.com>
---
 build/build.go     | 14 ++++++++++++--
 monitor/monitor.go | 16 ++++++++++++----
 2 files changed, 24 insertions(+), 6 deletions(-)

diff --git a/build/build.go b/build/build.go
index 5609377b..8974f0f5 100644
--- a/build/build.go
+++ b/build/build.go
@@ -643,7 +643,16 @@ func Invoke(ctx context.Context, cfg ContainerConfig) error {
 		return errors.Errorf("result must be provided")
 	}
 	c, res := cfg.ResultCtx.Client, cfg.ResultCtx.Res
-	_, err := c.Build(ctx, client.SolveOpt{}, "buildx", func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
+
+	mainCtx := ctx
+
+	_, err := c.Build(context.TODO(), client.SolveOpt{}, "buildx", func(ctx context.Context, c gateway.Client) (*gateway.Result, error) {
+		ctx, cancel := context.WithCancel(ctx)
+		go func() {
+			<-mainCtx.Done()
+			cancel()
+		}()
+
 		if res.Ref == nil {
 			return nil, errors.Errorf("no reference is registered")
 		}
@@ -673,7 +682,8 @@ func Invoke(ctx context.Context, cfg ContainerConfig) error {
 		if err != nil {
 			return nil, err
 		}
-		defer ctr.Release(ctx)
+		defer ctr.Release(context.TODO())
+
 		proc, err := ctr.Start(ctx, gateway.StartRequest{
 			Args:   cfg.Args,
 			Env:    cfg.Env,
diff --git a/monitor/monitor.go b/monitor/monitor.go
index 43688ed6..fbddbdd6 100644
--- a/monitor/monitor.go
+++ b/monitor/monitor.go
@@ -52,10 +52,8 @@ func RunMonitor(ctx context.Context, containerConfig build.ContainerConfig, relo
 	}
 
 	// Start container automatically
-	go func() {
-		fmt.Fprintf(stdout, "Launching interactive container. Press Ctrl-a-c to switch to monitor console\n")
-		m.rollback(ctx, containerConfig)
-	}()
+	fmt.Fprintf(stdout, "Launching interactive container. Press Ctrl-a-c to switch to monitor console\n")
+	m.rollback(ctx, containerConfig)
 
 	// Serve monitor commands
 	monitorForwarder := newIOForwarder(monitorIn)
@@ -67,6 +65,10 @@ func RunMonitor(ctx context.Context, containerConfig build.ContainerConfig, relo
 		go func() {
 			defer close(doneCh)
 			defer in.Close()
+			go func() {
+				<-ctx.Done()
+				in.Close()
+			}()
 			t := term.NewTerminal(readWriter{in.stdin, in.stdout}, "(buildx) ")
 			for {
 				l, err := t.ReadLine()
@@ -105,8 +107,14 @@ func RunMonitor(ctx context.Context, containerConfig build.ContainerConfig, relo
 		}()
 		select {
 		case <-doneCh:
+			if m.curInvokeCancel != nil {
+				m.curInvokeCancel()
+			}
 			return nil
 		case err := <-errCh:
+			if m.curInvokeCancel != nil {
+				m.curInvokeCancel()
+			}
 			return err
 		case <-monitorDisableCh:
 		}