Package ctxutil provides some helper functions for manipulating Golang contexts.

ctxutil.go

package ctxutil

import (
	"context"
	"os"
	"os/signal"
)

// An ErrSignal represents an error triggered by an operating system signal.
type ErrSignal struct {
	sig os.Signal
}

func (e *ErrSignal) Error() string {
	return "ctxutil: signal received: " + e.sig.String()
}

// WithSignal returns a copy of parent with a new Done channel. The new Done
// channel is closed when any of the given signals is received.
func WithSignal(parent context.Context, sig ...os.Signal) (context.Context, context.CancelFunc) {
	ch := make(chan os.Signal, 1)
	signal.Notify(ch, sig...)

	ctx, cancel := context.WithCancelCause(parent)
	go func() {
		select {
		case <-ctx.Done():
			cancel(nil)
		case sig := <-ch:
			cancel(&ErrSignal{sig})
		}
	}()

	return ctx, func() { cancel(nil) }
}

ctxutil_test.go

package ctxutil

import (
	"context"
	"os"
	"syscall"
	"testing"
	"time"
)

func TestWithSignal(t *testing.T) {
	parent := context.Background()
	ctx, cancel := WithSignal(parent, syscall.SIGINT)
	defer cancel()

	p, err := os.FindProcess(os.Getpid())
	if err != nil {
		t.Fatalf("failed finding current process: %s", err)
	}

	if err = p.Signal(syscall.SIGINT); err != nil {
		t.Fatalf("failed sending signal: %s", err)
	}

	// wait for a short time to make sure the channel has been closed
	time.Sleep(time.Millisecond)
	select {
	case <-ctx.Done():
		// success
	default:
		t.Error("WithSignal() failed")
	}
}