feat: Support graceful job step cancellation (#2714)

* feat: Support graceful job step cancellation

* for gh-act-runner
* act-cli support as well
* respecting always() and cancelled() of steps

* change main

* cancel startContainer / gh cli / bugreport early

* add to watch as well

---------

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

pkg/common/context.go

@@ -0,0 +1,45 @@
+package common
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+func createGracefulJobCancellationContext() (context.Context, func(), chan os.Signal) {
+	ctx := context.Background()
+	ctx, forceCancel := context.WithCancel(ctx)
+	cancelCtx, cancel := context.WithCancel(ctx)
+	ctx = WithJobCancelContext(ctx, cancelCtx)
+
+	// trap Ctrl+C and call cancel on the context
+	c := make(chan os.Signal, 1)
+	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
+	go func() {
+		select {
+		case sig := <-c:
+			if sig == os.Interrupt {
+				cancel()
+				select {
+				case <-c:
+					forceCancel()
+				case <-ctx.Done():
+				}
+			} else {
+				forceCancel()
+			}
+		case <-ctx.Done():
+		}
+	}()
+	return ctx, func() {
+		signal.Stop(c)
+		forceCancel()
+		cancel()
+	}, c
+}
+
+func CreateGracefulJobCancellationContext() (context.Context, func()) {
+	ctx, cancel, _ := createGracefulJobCancellationContext()
+	return ctx, cancel
+}

pkg/common/context_test.go

@@ -0,0 +1,98 @@
+package common
+
+import (
+	"context"
+	"os"
+	"syscall"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestGracefulJobCancellationViaSigint(t *testing.T) {
+	ctx, cancel, channel := createGracefulJobCancellationContext()
+	defer cancel()
+	assert.NotNil(t, ctx)
+	assert.NotNil(t, cancel)
+	assert.NotNil(t, channel)
+	cancelCtx := JobCancelContext(ctx)
+	assert.NotNil(t, cancelCtx)
+	assert.NoError(t, ctx.Err())
+	assert.NoError(t, cancelCtx.Err())
+	channel <- os.Interrupt
+	select {
+	case <-time.After(1 * time.Second):
+		t.Fatal("context not canceled")
+	case <-cancelCtx.Done():
+	case <-ctx.Done():
+	}
+	if assert.Error(t, cancelCtx.Err(), "context canceled") {
+		assert.Equal(t, context.Canceled, cancelCtx.Err())
+	}
+	assert.NoError(t, ctx.Err())
+	channel <- os.Interrupt
+	select {
+	case <-time.After(1 * time.Second):
+		t.Fatal("context not canceled")
+	case <-ctx.Done():
+	}
+	if assert.Error(t, ctx.Err(), "context canceled") {
+		assert.Equal(t, context.Canceled, ctx.Err())
+	}
+}
+
+func TestForceCancellationViaSigterm(t *testing.T) {
+	ctx, cancel, channel := createGracefulJobCancellationContext()
+	defer cancel()
+	assert.NotNil(t, ctx)
+	assert.NotNil(t, cancel)
+	assert.NotNil(t, channel)
+	cancelCtx := JobCancelContext(ctx)
+	assert.NotNil(t, cancelCtx)
+	assert.NoError(t, ctx.Err())
+	assert.NoError(t, cancelCtx.Err())
+	channel <- syscall.SIGTERM
+	select {
+	case <-time.After(1 * time.Second):
+		t.Fatal("context not canceled")
+	case <-cancelCtx.Done():
+	}
+	select {
+	case <-time.After(1 * time.Second):
+		t.Fatal("context not canceled")
+	case <-ctx.Done():
+	}
+	if assert.Error(t, ctx.Err(), "context canceled") {
+		assert.Equal(t, context.Canceled, ctx.Err())
+	}
+	if assert.Error(t, cancelCtx.Err(), "context canceled") {
+		assert.Equal(t, context.Canceled, cancelCtx.Err())
+	}
+}
+
+func TestCreateGracefulJobCancellationContext(t *testing.T) {
+	ctx, cancel := CreateGracefulJobCancellationContext()
+	defer cancel()
+	assert.NotNil(t, ctx)
+	assert.NotNil(t, cancel)
+	cancelCtx := JobCancelContext(ctx)
+	assert.NotNil(t, cancelCtx)
+	assert.NoError(t, cancelCtx.Err())
+}
+
+func TestCreateGracefulJobCancellationContextCancelFunc(t *testing.T) {
+	ctx, cancel := CreateGracefulJobCancellationContext()
+	assert.NotNil(t, ctx)
+	assert.NotNil(t, cancel)
+	cancelCtx := JobCancelContext(ctx)
+	assert.NotNil(t, cancelCtx)
+	assert.NoError(t, cancelCtx.Err())
+	cancel()
+	if assert.Error(t, ctx.Err(), "context canceled") {
+		assert.Equal(t, context.Canceled, ctx.Err())
+	}
+	if assert.Error(t, cancelCtx.Err(), "context canceled") {
+		assert.Equal(t, context.Canceled, cancelCtx.Err())
+	}
+}

pkg/common/job_error.go

@@ -8,6 +8,10 @@ type jobErrorContextKey string
 
 const jobErrorContextKeyVal = jobErrorContextKey("job.error")
 
+type jobCancelCtx string
+
+const JobCancelCtxVal = jobCancelCtx("job.cancel")
+
 // JobError returns the job error for current context if any
 func JobError(ctx context.Context) error {
 	val := ctx.Value(jobErrorContextKeyVal)
@@ -28,3 +32,35 @@ func WithJobErrorContainer(ctx context.Context) context.Context {
 	container := map[string]error{}
 	return context.WithValue(ctx, jobErrorContextKeyVal, container)
 }
+
+func WithJobCancelContext(ctx context.Context, cancelContext context.Context) context.Context {
+	return context.WithValue(ctx, JobCancelCtxVal, cancelContext)
+}
+
+func JobCancelContext(ctx context.Context) context.Context {
+	val := ctx.Value(JobCancelCtxVal)
+	if val != nil {
+		if container, ok := val.(context.Context); ok {
+			return container
+		}
+	}
+	return nil
+}
+
+// EarlyCancelContext returns a new context based on ctx that is canceled when the first of the provided contexts is canceled.
+func EarlyCancelContext(ctx context.Context) (context.Context, context.CancelFunc) {
+	val := JobCancelContext(ctx)
+	if val != nil {
+		context, cancel := context.WithCancel(ctx)
+		go func() {
+			defer cancel()
+			select {
+			case <-context.Done():
+			case <-ctx.Done():
+			case <-val.Done():
+			}
+		}()
+		return context, cancel
+	}
+	return ctx, func() {}
+}