[pkg/signaller] Add a package to handle Pause signal

This essentially is a copy of code used is xcontext
to do the same function. The major difference is that
now signals (like Pause) are completely decoupled from
context cancelling. Thus IsSignaledWith() and Until(nil)
do not catch a Context closure.

Signed-off-by: Dmitrii Okunev <[email protected]>

Author: xaionaro

pkg/signaller/broadcast_signal.go

@@ -0,0 +1,7 @@
+package signaller
+
+type broadcastSignalT struct{}
+
+func (broadcastSignalT) Error() string { return "broadcast-signal" }
+
+var broadcastSignal = broadcastSignalT{}

pkg/signaller/ctx_key.go

@@ -0,0 +1,5 @@
+package signaller
+
+type signallerCtxKeyT struct{}
+
+var signallerCtxKey = signallerCtxKeyT{}

pkg/signaller/is_signaled_with.go

@@ -0,0 +1,39 @@
+package signaller
+
+import (
+	"context"
+	"errors"
+)
+
+// IsSignaledWith returns true if the context received a signal equals
+// to any of passed ones.
+//
+// If signals is empty, then returns true if the context received any
+// signal.
+func IsSignaledWith(ctx context.Context, signals ...error) bool {
+	s := getSignaller(ctx)
+	if s == nil {
+		return false
+	}
+
+	s.locker.Lock()
+	defer s.locker.Unlock()
+
+	return s.isSignaledWith(signals...)
+}
+
+func (s *signaller) isSignaledWith(signals ...error) bool {
+	if len(signals) == 0 && (len(s.receivedSignals) != 0) {
+		return true
+	}
+
+	for _, receivedErr := range s.receivedSignals {
+		for _, err := range signals {
+			if errors.Is(receivedErr, err) {
+				return true
+			}
+		}
+	}
+
+	return false
+}

pkg/signaller/received_signals.go

@@ -0,0 +1,23 @@
+package signaller
+
+import "context"
+
+func ReceivedSignals(ctx context.Context) []error {
+	s := getSignaller(ctx)
+	return s.ReceivedSignals()
+}
+
+// ReceivedSignals returns all the received signals (including events
+// received by parents).
+//
+// This is a read-only value, do not modify it.
+func (s *signaller) ReceivedSignals() []error {
+	if s == nil {
+		return nil
+	}
+
+	s.locker.Lock()
+	defer s.locker.Unlock()
+
+	return s.receivedSignals
+}

pkg/signaller/signaller.go

@@ -0,0 +1,64 @@
+package signaller
+
+import (
+	"context"
+	"runtime"
+	"sync"
+)
+
+var (
+	openChan   = make(chan struct{})
+	closedChan = make(chan struct{})
+)
+
+func init() {
+	close(closedChan)
+}
+
+type signaller struct {
+	// locker is used exclude concurrent access to any data below (in this
+	// structure).
+	locker sync.Mutex
+
+	// children are all signallers derived from this one. And for example
+	// if this one will a signal it will be propagated to all the children.
+	children map[*signaller]struct{}
+
+	// receivedSignals accumulates any signals ever received by this instance
+	// and/or by any parent.
+	receivedSignals []error
+
+	// signalChans are the channels returned by Until.
+	signalChans map[error]chan struct{}
+}
+
+func getSignaller(ctx context.Context) *signaller {
+	signaller, _ := ctx.Value(signallerCtxKey).(*signaller)
+	return signaller
+}
+
+func withSignaller(ctx context.Context) (*signaller, context.Context) {
+	parent := getSignaller(ctx)
+	s := &signaller{
+		children:    make(map[*signaller]struct{}),
+		signalChans: make(map[error]chan struct{}),
+	}
+	ctx = context.WithValue(ctx, signallerCtxKey, s)
+	if parent == nil {
+		return s, ctx
+	}
+
+	parent.locker.Lock()
+	s.receivedSignals = make([]error, len(parent.receivedSignals))
+	copy(s.receivedSignals, parent.receivedSignals)
+	parent.children[s] = struct{}{}
+	parent.locker.Unlock()
+
+	runtime.SetFinalizer(s, func(s *signaller) {
+		parent.locker.Lock()
+		delete(parent.children, s)
+		parent.locker.Unlock()
+	})
+
+	return s, ctx
+}

pkg/signaller/signaller_test.go

@@ -0,0 +1,118 @@
+// Copyright (c) Facebook, Inc. and its affiliates.
+//
+// This source code is licensed under the MIT license found in the
+// LICENSE file in the root directory of this source tree.
+
+package signaller
+
+import (
+	"context"
+	"fmt"
+	"sync"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+var unitTestCustomSignal0 = fmt.Errorf("unit-test custom signal #0")
+var unitTestCustomSignal1 = fmt.Errorf("unit-test custom signal #1")
+var unitTestCustomSignal2 = fmt.Errorf("unit-test custom signal #2")
+
+func TestBackgroundContext(t *testing.T) {
+	ctx := context.Background()
+
+	var blocked bool
+	select {
+	case <-Until(ctx, nil):
+	default:
+		blocked = true
+	}
+
+	require.True(t, blocked)
+	require.Nil(t, ReceivedSignals(ctx))
+}
+
+func TestRace(t *testing.T) {
+	ctx, sendSignal := WithSignal(context.Background(), unitTestCustomSignal0)
+	var wg sync.WaitGroup
+	readyToSelect := make(chan struct{})
+	startSelect := make(chan struct{})
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		ch := Until(ctx, nil)
+		select {
+		case <-ch:
+			t.Fail()
+		default:
+		}
+		close(readyToSelect)
+		<-startSelect
+		<-ch
+	}()
+	<-readyToSelect
+	sendSignal()
+	close(startSelect)
+	wg.Wait()
+}
+
+func TestSendSignal(t *testing.T) {
+	ctx, pauseFunc := WithSignal(context.Background(), unitTestCustomSignal0)
+	require.NotNil(t, ctx)
+
+	pauseFunc()
+
+	<-Until(ctx, unitTestCustomSignal0)
+	<-Until(ctx, nil)
+
+	require.Nil(t, ctx.Err())
+	require.Equal(t, []error{unitTestCustomSignal0}, ReceivedSignals(ctx))
+
+	var canceled bool
+	select {
+	case <-ctx.Done():
+		canceled = true
+	default:
+	}
+	require.False(t, canceled)
+}
+
+func TestSendMultipleSignals(t *testing.T) {
+	ctx := context.Background()
+	ctx, sendSignal0 := WithSignal(ctx, unitTestCustomSignal0)
+	ctx, sendSignal1 := WithSignal(ctx, unitTestCustomSignal1)
+	require.NotNil(t, ctx)
+
+	sendSignal1()
+	sendSignal0()
+	sendSignal1()
+	sendSignal0()
+
+	<-Until(ctx, nil)
+	<-Until(ctx, unitTestCustomSignal0)
+	<-Until(ctx, unitTestCustomSignal1)
+
+	require.Equal(t, []error{unitTestCustomSignal1, unitTestCustomSignal0, unitTestCustomSignal1, unitTestCustomSignal0}, ReceivedSignals(ctx))
+}
+
+func TestGrandGrandGrandChild(t *testing.T) {
+	type myUniqueType string
+	ctx0, sendSignal0 := WithSignal(context.Background(), unitTestCustomSignal0)
+	ctx1, _ := WithSignal(context.WithValue(ctx0, myUniqueType("someKey1"), "someValue1"), unitTestCustomSignal1)
+	ctx2, _ := WithSignal(context.WithValue(ctx1, myUniqueType("someKey2"), "someValue2"), unitTestCustomSignal2)
+
+	require.False(t, IsSignaledWith(ctx2, unitTestCustomSignal0))
+	sendSignal0()
+	<-Until(ctx2, unitTestCustomSignal0)
+	require.True(t, IsSignaledWith(ctx2, unitTestCustomSignal0))
+	require.False(t, IsSignaledWith(ctx2, unitTestCustomSignal1))
+
+	select {
+	case <-Until(ctx2, unitTestCustomSignal1):
+		require.FailNow(t, "unexpected closed chan for signal1")
+	case <-Until(ctx2, unitTestCustomSignal2):
+		require.FailNow(t, "unexpected closed chan for signal2")
+	default:
+	}
+
+}

pkg/signaller/until.go

@@ -0,0 +1,35 @@
+package signaller
+
+import (
+	"context"
+)
+
+// Until works similar to Done(), but it is possible to specify specific
+// signal to wait for.
+//
+// If signal is nil, then waits for any event.
+func Until(ctx context.Context, signal error) <-chan struct{} {
+	s := getSignaller(ctx)
+	if s == nil {
+		return openChan
+	}
+	return s.getSignalChan(signal)
+}
+
+func (s *signaller) getSignalChan(signal error) <-chan struct{} {
+	if signal == nil {
+		signal = broadcastSignal
+	}
+
+	s.locker.Lock()
+	defer s.locker.Unlock()
+
+	if (signal == broadcastSignal && s.isSignaledWith()) || (signal != broadcastSignal && s.isSignaledWith(signal)) {
+		return closedChan
+	}
+
+	if s.signalChans[signal] == nil {
+		s.signalChans[signal] = make(chan struct{})
+	}
+	return s.signalChans[signal]
+}

pkg/signaller/with_signal.go

@@ -0,0 +1,45 @@
+package signaller
+
+import (
+	"context"
+)
+
+// WithSignal is similar to context.WithCancel, but it does not close the context
+// and sends a signal to the Context which use observable using Until and IsSignaledWith
+// functions.
+func WithSignal(ctx context.Context, signal error, moreSignals ...error) (context.Context, context.CancelFunc) {
+	s, ctx := withSignaller(ctx)
+	return ctx, s.signalFunc(signal, moreSignals...)
+}
+
+func (s *signaller) signalFunc(signal error, moreSignals ...error) context.CancelFunc {
+	return func() {
+		s.signal(signal, moreSignals...)
+	}
+}
+
+func (s *signaller) signal(signal error, moreSignals ...error) {
+	s.locker.Lock()
+	defer s.locker.Unlock()
+
+	s.receivedSignals = append(s.receivedSignals, signal)
+	s.receivedSignals = append(s.receivedSignals, moreSignals...)
+
+	s.sendOneSignal(signal)
+	for _, signal := range moreSignals {
+		s.sendOneSignal(signal)
+	}
+	s.sendOneSignal(broadcastSignal)
+
+	for child := range s.children {
+		child.signal(signal, moreSignals...)
+	}
+}
+
+func (s *signaller) sendOneSignal(signal error) {
+	if s.signalChans[signal] == nil {
+		return
+	}
+	close(s.signalChans[signal])
+	s.signalChans[signal] = nil
+}

pkg/signaller/without_signaller.go

@@ -0,0 +1,10 @@
+package signaller
+
+import (
+	"context"
+)
+
+// WithoutSignaller derives a context without any signals setup.
+func WithoutSignaller(ctx context.Context) context.Context {
+	return context.WithValue(ctx, signallerCtxKey, nil)
+}

pkg/signals/paused.go

@@ -0,0 +1,7 @@
+package signals
+
+type pausedT struct{}
+
+func (pausedT) Error() string { return "job is paused" }
+
+var Paused = error(pausedT{})