feat: use signal-wrapped contexts in CLI commands (#38)

pkg/signals/context.go

@@ -0,0 +1,31 @@
+package signals
+
+import (
+	"context"
+	"os"
+	"os/signal"
+	"syscall"
+)
+
+// WithSignals returns a context that is canceled with any signal in sigs.
+func WithSignals(ctx context.Context, sigs ...os.Signal) context.Context {
+	sigCh := make(chan os.Signal, 1)
+	signal.Notify(sigCh, sigs...)
+
+	ctx, cancel := context.WithCancel(ctx)
+	go func() {
+		defer cancel()
+		select {
+		case <-ctx.Done():
+			return
+		case <-sigCh:
+			return
+		}
+	}()
+	return ctx
+}
+
+// WithStandardSignals cancels the context on os.Interrupt, syscall.SIGTERM.
+func WithStandardSignals(ctx context.Context) context.Context {
+	return WithSignals(ctx, os.Interrupt, syscall.SIGTERM)
+}

pkg/signals/context_test.go

@@ -0,0 +1,79 @@
+package signals
+
+import (
+	"context"
+	"fmt"
+	"os"
+	"syscall"
+	"testing"
+	"time"
+)
+
+func ExampleWithSignals() {
+	ctx := WithSignals(context.Background(), syscall.SIGUSR1)
+	go func() {
+		time.Sleep(500 * time.Millisecond) // after some time SIGUSR1 is sent
+		// mimicking a signal from the outside
+		syscall.Kill(syscall.Getpid(), syscall.SIGUSR1)
+	}()
+
+	<-ctx.Done()
+	fmt.Println("finished")
+	// Output:
+	// finished
+}
+
+func Example_withUnregisteredSignals() {
+	dctx, cancel := context.WithTimeout(context.TODO(), time.Millisecond*100)
+	defer cancel()
+
+	ctx := WithSignals(dctx, syscall.SIGUSR1)
+	go func() {
+		time.Sleep(10 * time.Millisecond) // after some time SIGUSR2 is sent
+		// mimicking a signal from the outside, WithSignals will not handle it
+		syscall.Kill(syscall.Getpid(), syscall.SIGUSR2)
+	}()
+
+	<-ctx.Done()
+	fmt.Println("finished")
+	// Output:
+	// finished
+}
+
+func TestWithSignals(t *testing.T) {
+	tests := []struct {
+		name       string
+		ctx        context.Context
+		sigs       []os.Signal
+		wantSignal bool
+	}{
+		{
+			name:       "sending signal SIGUSR2 should exit context.",
+			ctx:        context.Background(),
+			sigs:       []os.Signal{syscall.SIGUSR2},
+			wantSignal: true,
+		},
+		{
+			name: "sending signal SIGUSR2 should NOT exit context.",
+			ctx:  context.Background(),
+			sigs: []os.Signal{syscall.SIGUSR1},
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			ctx := WithSignals(tt.ctx, tt.sigs...)
+			syscall.Kill(syscall.Getpid(), syscall.SIGUSR2)
+			timer := time.NewTimer(500 * time.Millisecond)
+			select {
+			case <-ctx.Done():
+				if !tt.wantSignal {
+					t.Errorf("unexpected exit with signal")
+				}
+			case <-timer.C:
+				if tt.wantSignal {
+					t.Errorf("expected to exit with signal but did not")
+				}
+			}
+		})
+	}
+}