Graceful cancellation for network dependent commands

cmd/ack-generate/command/apis.go

@@ -14,6 +14,7 @@
 package command
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"path/filepath"
@@ -91,7 +92,9 @@ func generateAPIs(cmd *cobra.Command, args []string) error {
 	if optOutputPath == "" {
 		optOutputPath = filepath.Join(optServicesDir, svcAlias)
 	}
-	if err := ensureSDKRepo(optCacheDir, optRefreshCache); err != nil {
+	ctx, cancel := contextWithSigterm(context.Background())
+	defer cancel()
+	if err := ensureSDKRepo(ctx, optCacheDir, optRefreshCache); err != nil {
 		return err
 	}
 	sdkHelper := model.NewSDKHelper(sdkDir)

cmd/ack-generate/command/common.go

@@ -18,18 +18,46 @@ import (
 	"fmt"
 	"io/ioutil"
 	"os"
+	"os/signal"
 	"path/filepath"
 	"strings"
+	"syscall"
+	"time"
 
 	"golang.org/x/mod/modfile"
 
 	"github.com/aws-controllers-k8s/code-generator/pkg/util"
 )
 
 const (
-	sdkRepoURL = "https://github.com/aws/aws-sdk-go"
+	sdkRepoURL             = "https://github.com/aws/aws-sdk-go"
+	defaultGitCloneTimeout = 180 * time.Second
+	defaultGitFetchTimeout = 30 * time.Second
 )
 
+func contextWithSigterm(ctx context.Context) (context.Context, context.CancelFunc) {
+	ctx, cancel := context.WithCancel(ctx)
+	signalCh := make(chan os.Signal, 1)
+
+	// recreate the context.CancelFunc
+	cancelFunc := func() {
+		signal.Stop(signalCh)
+		cancel()
+	}
+
+	// notify on SIGINT or SIGTERM
+	signal.Notify(signalCh, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		select {
+		case <-signalCh:
+			cancel()
+		case <-ctx.Done():
+		}
+	}()
+
+	return ctx, cancelFunc
+}
+
 // ensureDir makes sure that a supplied directory exists and
 // returns whether the directory already existed.
 func ensureDir(fp string) (bool, error) {
@@ -69,6 +97,7 @@ func isDirWriteable(fp string) bool {
 // the aws-sdk-go is found. It will also optionally fetch all the remote tags
 // and checkout the given tag.
 func ensureSDKRepo(
+	ctx context.Context,
 	cacheDir string,
 	// A boolean instructing ensureSDKRepo whether to fetch the remote tags from
 	// the upstream repository
@@ -83,15 +112,20 @@ func ensureSDKRepo(
 	// Clone repository if it doen't exist
 	sdkDir = filepath.Join(srcPath, "aws-sdk-go")
 	if _, err := os.Stat(sdkDir); os.IsNotExist(err) {
-		err = util.CloneRepository(context.Background(), sdkDir, sdkRepoURL)
+
+		ctx, cancel := context.WithTimeout(ctx, defaultGitCloneTimeout)
+		defer cancel()
+		err = util.CloneRepository(ctx, sdkDir, sdkRepoURL)
 		if err != nil {
 			return fmt.Errorf("canot clone repository: %v", err)
 		}
 	}
 
 	// Fetch all tags
 	if fetchTags {
-		err = util.FetchRepositoryTags(context.Background(), sdkDir)
+		ctx, cancel := context.WithTimeout(ctx, defaultGitFetchTimeout)
+		defer cancel()
+		err = util.FetchRepositoryTags(ctx, sdkDir)
 		if err != nil {
 			return fmt.Errorf("cannot fetch tags: %v", err)
 		}

cmd/ack-generate/command/controller.go

@@ -15,6 +15,7 @@ package command
 
 import (
 	"bufio"
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -57,7 +58,9 @@ func generateController(cmd *cobra.Command, args []string) error {
 		optOutputPath = filepath.Join(optServicesDir, svcAlias)
 	}
 
-	if err := ensureSDKRepo(optCacheDir, optRefreshCache); err != nil {
+	ctx, cancel := contextWithSigterm(context.Background())
+	defer cancel()
+	if err := ensureSDKRepo(ctx, optCacheDir, optRefreshCache); err != nil {
 		return err
 	}
 	sdkHelper := ackmodel.NewSDKHelper(sdkDir)

cmd/ack-generate/command/crossplane.go

@@ -14,6 +14,7 @@
 package command
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"os"
@@ -49,7 +50,9 @@ func generateCrossplane(_ *cobra.Command, args []string) error {
 	if len(args) != 1 {
 		return fmt.Errorf("please specify the service alias for the AWS service API to generate")
 	}
-	if err := ensureSDKRepo(optCacheDir, optRefreshCache); err != nil {
+	ctx, cancel := contextWithSigterm(context.Background())
+	defer cancel()
+	if err := ensureSDKRepo(ctx, optCacheDir, optRefreshCache); err != nil {
 		return err
 	}
 	svcAlias := strings.ToLower(args[0])

cmd/ack-generate/command/olm.go

@@ -14,6 +14,7 @@
 package command
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"path/filepath"
@@ -81,7 +82,9 @@ func generateOLMAssets(cmd *cobra.Command, args []string) error {
 	version := args[1]
 
 	// get the generator inputs
-	if err := ensureSDKRepo(optCacheDir, optRefreshCache); err != nil {
+	ctx, cancel := contextWithSigterm(context.Background())
+	defer cancel()
+	if err := ensureSDKRepo(ctx, optCacheDir, optRefreshCache); err != nil {
 		return err
 	}
 	sdkHelper := ackmodel.NewSDKHelper(sdkDir)

cmd/ack-generate/command/release.go

@@ -14,6 +14,7 @@
 package command
 
 import (
+	"context"
 	"fmt"
 	"io/ioutil"
 	"path/filepath"
@@ -65,7 +66,9 @@ func generateRelease(cmd *cobra.Command, args []string) error {
 	// version supplied hasn't been used (as a Git tag) before...
 	releaseVersion := strings.ToLower(args[1])
 
-	if err := ensureSDKRepo(optCacheDir, optRefreshCache); err != nil {
+	ctx, cancel := contextWithSigterm(context.Background())
+	defer cancel()
+	if err := ensureSDKRepo(ctx, optCacheDir, optRefreshCache); err != nil {
 		return err
 	}
 	sdkHelper := ackmodel.NewSDKHelper(sdkDir)