Package ctxutil provides some helper functions for manipulating Golang contexts.

ctxutil.go

package ctxutil

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

// 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) }
}

// Shorten makes the deadline of the parent context shorter by the given duration.
// If parent does not have a deadline, Shorten returns a copy of the context with
// a new Done channel.
func Shorten(parent context.Context, dur time.Duration) (context.Context, context.CancelFunc) {
	if deadline, ok := parent.Deadline(); ok {
		return context.WithDeadline(parent, deadline.Add(-dur))
	}
	return context.WithCancel(parent)
}

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")
	}
}

func TestShorten(t *testing.T) {
	parent, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	dur := 500 * time.Millisecond
	ctx, shortenCancel := Shorten(parent, dur)
	defer shortenCancel()

	parentDeadline, ok := parent.Deadline()
	if !ok {
		t.Fatal("parent context has no deadline")
	}
	deadline, ok := ctx.Deadline()
	if !ok {
		t.Fatal("Shorten() returned a context without a deadline")
	}
	if diff := parentDeadline.Sub(deadline); diff != dur {
		t.Errorf("Shorten() deadline diff = %s, want %s", diff, dur)
	}

	select {
	case <-ctx.Done():
		t.Error("Shorten() context done too early")
	default:
	}

	time.Sleep(dur + 50*time.Millisecond)
	select {
	case <-ctx.Done():
		// success
	default:
		t.Error("Shorten() context not done after shortened deadline")
	}
}