Add context to go build (ko-build#105)

pkg/build/build.go

@@ -15,6 +15,8 @@
 package build
 
 import (
+	"context"
+
 	v1 "github.com/google/go-containerregistry/pkg/v1"
 )
 
@@ -27,5 +29,5 @@ type Interface interface {
 	IsSupportedReference(string) bool
 
 	// Build turns the given importpath reference into a v1.Image containing the Go binary.
-	Build(string) (v1.Image, error)
+	Build(context.Context, string) (v1.Image, error)
 }

pkg/build/gobuild.go

@@ -18,6 +18,7 @@ import (
 	"archive/tar"
 	"bytes"
 	"compress/gzip"
+	"context"
 	"encoding/json"
 	"errors"
 	gb "go/build"
@@ -41,7 +42,7 @@ const (
 
 // GetBase takes an importpath and returns a base v1.Image.
 type GetBase func(string) (v1.Image, error)
-type builder func(string, v1.Platform, bool) (string, error)
+type builder func(context.Context, string, v1.Platform, bool) (string, error)
 
 type gobuild struct {
 	getBase              GetBase
@@ -147,7 +148,7 @@ func (g *gobuild) importPackage(s string) (*gb.Package, error) {
 	return nil, moduleErr
 }
 
-func build(ip string, platform v1.Platform, disableOptimizations bool) (string, error) {
+func build(ctx context.Context, ip string, platform v1.Platform, disableOptimizations bool) (string, error) {
 	tmpDir, err := ioutil.TempDir("", "ko")
 	if err != nil {
 		return "", err
@@ -163,7 +164,7 @@ func build(ip string, platform v1.Platform, disableOptimizations bool) (string,
 	args = append(args, "-o", file)
 	args = addGo113TrimPathFlag(args)
 	args = append(args, ip)
-	cmd := exec.Command("go", args...)
+	cmd := exec.CommandContext(ctx, "go", args...)
 
 	// Last one wins
 	defaultEnv := []string{
@@ -366,7 +367,7 @@ func (g *gobuild) tarKoData(importpath string) (*bytes.Buffer, error) {
 }
 
 // Build implements build.Interface
-func (gb *gobuild) Build(s string) (v1.Image, error) {
+func (gb *gobuild) Build(ctx context.Context, s string) (v1.Image, error) {
 	// Determine the appropriate base image for this import path.
 	base, err := gb.getBase(s)
 	if err != nil {
@@ -382,7 +383,7 @@ func (gb *gobuild) Build(s string) (v1.Image, error) {
 	}
 
 	// Do the build into a temporary file.
-	file, err := gb.build(s, platform, gb.disableOptimizations)
+	file, err := gb.build(ctx, s, platform, gb.disableOptimizations)
 	if err != nil {
 		return nil, err
 	}

pkg/build/gobuild_test.go

@@ -16,6 +16,7 @@ package build
 
 import (
 	"archive/tar"
+	"context"
 	"io"
 	"io/ioutil"
 	"path/filepath"
@@ -102,7 +103,7 @@ func TestGoBuildIsSupportedRefWithModules(t *testing.T) {
 }
 
 // A helper method we use to substitute for the default "build" method.
-func writeTempFile(s string, _ v1.Platform, _ bool) (string, error) {
+func writeTempFile(_ context.Context, s string, _ v1.Platform, _ bool) (string, error) {
 	tmpDir, err := ioutil.TempDir("", "ko")
 	if err != nil {
 		return "", err
@@ -137,7 +138,7 @@ func TestGoBuildNoKoData(t *testing.T) {
 		t.Fatalf("NewGo() = %v", err)
 	}
 
-	img, err := ng.Build(filepath.Join(importpath, "cmd", "ko"))
+	img, err := ng.Build(context.Background(), filepath.Join(importpath, "cmd", "ko"))
 	if err != nil {
 		t.Fatalf("Build() = %v", err)
 	}
@@ -217,7 +218,7 @@ func TestGoBuild(t *testing.T) {
 		t.Fatalf("NewGo() = %v", err)
 	}
 
-	img, err := ng.Build(filepath.Join(importpath, "cmd", "ko", "test"))
+	img, err := ng.Build(context.Background(), filepath.Join(importpath, "cmd", "ko", "test"))
 	if err != nil {
 		t.Fatalf("Build() = %v", err)
 	}

pkg/build/limit.go

@@ -36,14 +36,14 @@ func (l *Limiter) IsSupportedReference(ip string) bool {
 }
 
 // Build implements Interface
-func (l *Limiter) Build(ip string) (v1.Image, error) {
+func (l *Limiter) Build(ctx context.Context, ip string) (v1.Image, error) {
 	// TODO(jonjohnsonjr): Build should take a context.Context.
 	if err := l.semaphore.Acquire(context.TODO(), 1); err != nil {
 		return nil, err
 	}
 	defer l.semaphore.Release(1)
 
-	return l.Builder.Build(ip)
+	return l.Builder.Build(ctx, ip)
 }
 
 // NewLimiter returns a new builder that only allows n concurrent builds of b.

pkg/build/limit_test.go

@@ -33,7 +33,7 @@ func (r *sleeper) IsSupportedReference(ip string) bool {
 }
 
 // Build implements Interface
-func (r *sleeper) Build(ip string) (v1.Image, error) {
+func (r *sleeper) Build(_ context.Context, ip string) (v1.Image, error) {
 	time.Sleep(50 * time.Millisecond)
 	return nil, nil
 }
@@ -45,7 +45,7 @@ func TestLimiter(t *testing.T) {
 	g, _ := errgroup.WithContext(context.TODO())
 	for i := 0; i <= 10; i++ {
 		g.Go(func() error {
-			_, _ = b.Build("whatever")
+			_, _ = b.Build(context.Background(), "whatever")
 			return nil
 		})
 	}

pkg/build/recorder.go

@@ -15,6 +15,7 @@
 package build
 
 import (
+	"context"
 	"sync"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -36,11 +37,11 @@ func (r *Recorder) IsSupportedReference(ip string) bool {
 }
 
 // Build implements Interface
-func (r *Recorder) Build(ip string) (v1.Image, error) {
+func (r *Recorder) Build(ctx context.Context, ip string) (v1.Image, error) {
 	func() {
 		r.m.Lock()
 		defer r.m.Unlock()
 		r.ImportPaths = append(r.ImportPaths, ip)
 	}()
-	return r.Builder.Build(ip)
+	return r.Builder.Build(ctx, ip)
 }

pkg/build/recorder_test.go

@@ -15,6 +15,7 @@
 package build
 
 import (
+	"context"
 	"testing"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -35,7 +36,7 @@ func (r *fake) IsSupportedReference(ip string) bool {
 }
 
 // Build implements Interface
-func (r *fake) Build(ip string) (v1.Image, error) {
+func (r *fake) Build(_ context.Context, ip string) (v1.Image, error) {
 	return r.b(ip)
 }
 
@@ -111,7 +112,7 @@ func TestBuildRecording(t *testing.T) {
 				Builder: inner,
 			}
 			for _, in := range test.inputs {
-				rec.Build(in)
+				rec.Build(context.Background(), in)
 			}
 			if diff := cmp.Diff(test.inputs, rec.ImportPaths); diff != "" {
 				t.Errorf("Build (-want, +got): %s", diff)

pkg/build/shared.go

@@ -15,6 +15,7 @@
 package build
 
 import (
+	"context"
 	"sync"
 
 	v1 "github.com/google/go-containerregistry/pkg/v1"
@@ -43,7 +44,7 @@ func NewCaching(inner Interface) (*Caching, error) {
 }
 
 // Build implements Interface
-func (c *Caching) Build(ip string) (v1.Image, error) {
+func (c *Caching) Build(ctx context.Context, ip string) (v1.Image, error) {
 	f := func() *future {
 		// Lock the map of futures.
 		c.m.Lock()
@@ -56,7 +57,7 @@ func (c *Caching) Build(ip string) (v1.Image, error) {
 		}
 		// Otherwise create and record a future for a Build of "ip".
 		f = newFuture(func() (v1.Image, error) {
-			return c.inner.Build(ip)
+			return c.inner.Build(ctx, ip)
 		})
 		c.results[ip] = f
 		return f

pkg/build/shared_test.go

@@ -15,6 +15,7 @@
 package build
 
 import (
+	"context"
 	"testing"
 	"time"
 
@@ -33,7 +34,7 @@ func (sb *slowbuild) IsSupportedReference(string) bool {
 	return true
 }
 
-func (sb *slowbuild) Build(string) (v1.Image, error) {
+func (sb *slowbuild) Build(_ context.Context, _ string) (v1.Image, error) {
 	time.Sleep(sb.sleep)
 	return random.Image(256, 8)
 }
@@ -55,7 +56,7 @@ func TestCaching(t *testing.T) {
 	// cache and iterate.
 	for idx := 0; idx < 3; idx++ {
 		start := time.Now()
-		img1, err := cb.Build(ip)
+		img1, err := cb.Build(context.Background(), ip)
 		if err != nil {
 			t.Errorf("Build() = %v", err)
 		}
@@ -73,7 +74,7 @@ func TestCaching(t *testing.T) {
 		previousDigest = d1
 
 		start = time.Now()
-		img2, err := cb.Build(ip)
+		img2, err := cb.Build(context.Background(), ip)
 		if err != nil {
 			t.Errorf("Build() = %v", err)
 		}

pkg/commands/apply.go

@@ -117,7 +117,8 @@ func addApply(topLevel *cobra.Command) {
 					stdin.Write([]byte("---\n"))
 				}
 				// Once primed kick things off.
-				resolveFilesToWriter(builder, publisher, fo, so, sto, stdin)
+				ctx := createCancellableContext()
+				resolveFilesToWriter(ctx, builder, publisher, fo, so, sto, stdin)
 			}()
 
 			// Run it.

pkg/commands/config.go

@@ -15,10 +15,13 @@
 package commands
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"os"
+	"os/signal"
 	"strconv"
+	"syscall"
 	"time"
 
 	"github.com/google/go-containerregistry/pkg/authn"
@@ -55,6 +58,19 @@ func getCreationTime() (*v1.Time, error) {
 	return &v1.Time{time.Unix(seconds, 0)}, nil
 }
 
+func createCancellableContext() context.Context {
+	signals := make(chan os.Signal)
+	signal.Notify(signals, syscall.SIGINT, syscall.SIGTERM)
+	ctx, cancel := context.WithCancel(context.Background())
+
+	go func() {
+		<-signals
+		cancel()
+	}()
+
+	return ctx
+}
+
 func init() {
 	// If omitted, use this base image.
 	viper.SetDefault("defaultBaseImage", "gcr.io/distroless/static:latest")

pkg/commands/create.go

@@ -117,7 +117,8 @@ func addCreate(topLevel *cobra.Command) {
 					stdin.Write([]byte("---\n"))
 				}
 				// Once primed kick things off.
-				resolveFilesToWriter(builder, publisher, fo, so, sto, stdin)
+				ctx := createCancellableContext()
+				resolveFilesToWriter(ctx, builder, publisher, fo, so, sto, stdin)
 			}()
 
 			// Run it.

pkg/commands/publish.go

@@ -68,7 +68,8 @@ func addPublish(topLevel *cobra.Command) {
 			if err != nil {
 				log.Fatalf("error creating publisher: %v", err)
 			}
-			images, err := publishImages(args, publisher, builder)
+			ctx := createCancellableContext()
+			images, err := publishImages(ctx, args, publisher, builder)
 			if err != nil {
 				log.Fatalf("failed to publish images: %v", err)
 			}

pkg/commands/publisher.go

@@ -15,6 +15,7 @@
 package commands
 
 import (
+	"context"
 	"fmt"
 	gb "go/build"
 
@@ -39,7 +40,7 @@ func qualifyLocalImport(importpath string) (string, error) {
 	return pkgs[0].PkgPath, nil
 }
 
-func publishImages(importpaths []string, pub publish.Interface, b build.Interface) (map[string]name.Reference, error) {
+func publishImages(ctx context.Context, importpaths []string, pub publish.Interface, b build.Interface) (map[string]name.Reference, error) {
 	imgs := make(map[string]name.Reference)
 	for _, importpath := range importpaths {
 		if gb.IsLocalImport(importpath) {
@@ -54,7 +55,7 @@ func publishImages(importpaths []string, pub publish.Interface, b build.Interfac
 			return nil, fmt.Errorf("importpath %q is not supported", importpath)
 		}
 
-		img, err := b.Build(importpath)
+		img, err := b.Build(ctx, importpath)
 		if err != nil {
 			return nil, fmt.Errorf("error building %q: %v", importpath, err)
 		}

pkg/commands/resolve.go

@@ -66,7 +66,8 @@ func addResolve(topLevel *cobra.Command) {
 			if err != nil {
 				log.Fatalf("error creating publisher: %v", err)
 			}
-			resolveFilesToWriter(builder, publisher, fo, so, sto, os.Stdout)
+			ctx := createCancellableContext()
+			resolveFilesToWriter(ctx, builder, publisher, fo, so, sto, os.Stdout)
 		},
 	}
 	options.AddLocalArg(resolve, lo)

pkg/commands/resolver.go

@@ -15,6 +15,7 @@
 package commands
 
 import (
+	"context"
 	"bytes"
 	"errors"
 	"fmt"
@@ -120,6 +121,7 @@ func makePublisher(no *options.NameOptions, lo *options.LocalOptions, ta *option
 type resolvedFuture chan []byte
 
 func resolveFilesToWriter(
+	ctx context.Context, 
 	builder *build.Caching,
 	publisher publish.Interface,
 	fo *options.FilenameOptions,
@@ -206,7 +208,7 @@ func resolveFilesToWriter(
 				recordingBuilder := &build.Recorder{
 					Builder: builder,
 				}
-				b, err := resolveFile(f, recordingBuilder, publisher, so, sto)
+				b, err := resolveFile(ctx, f, recordingBuilder, publisher, so, sto)
 				if err != nil {
 					// Don't let build errors disrupt the watch.
 					lg := log.Fatalf
@@ -252,6 +254,7 @@ func resolveFilesToWriter(
 }
 
 func resolveFile(
+	ctx context.Context,
 	f string,
 	builder build.Interface,
 	pub publish.Interface,
@@ -304,7 +307,7 @@ func resolveFile(
 
 	}
 
-	if err := resolve.ImageReferences(docNodes, sto.Strict, builder, pub); err != nil {
+	if err := resolve.ImageReferences(ctx, docNodes, sto.Strict, builder, pub); err != nil {
 		return nil, fmt.Errorf("error resolving image references: %v", err)
 	}