move implementation to internal/ directory (#828)

internal/util/optionaldeadline/optionaldeadline.go

@@ -0,0 +1,83 @@
+package optionaldeadline
+
+import (
+	"context"
+	"sync"
+	"time"
+)
+
+type contextWithOptionalDeadline struct {
+	context.Context
+
+	m        sync.Mutex
+	deadline time.Time
+
+	done chan struct{}
+	err  error
+}
+
+func (c *contextWithOptionalDeadline) Deadline() (deadline time.Time, ok bool) {
+	c.m.Lock()
+	defer c.m.Unlock()
+	return c.deadline, !c.deadline.IsZero()
+}
+
+func (c *contextWithOptionalDeadline) Err() error {
+	c.m.Lock()
+	defer c.m.Unlock()
+	return c.err
+}
+
+func (c *contextWithOptionalDeadline) Done() <-chan struct{} {
+	return c.done
+}
+
+func ContextWithOptionalDeadline(pctx context.Context) (ctx context.Context, enforceDeadline func(deadline time.Time)) {
+
+	// mctx can only be cancelled by cancelMctx, not by a potential cancel of pctx
+	rctx := &contextWithOptionalDeadline{
+		Context: pctx,
+		done:    make(chan struct{}),
+		err:     nil,
+	}
+	enforceDeadline = func(deadline time.Time) {
+
+		// Set deadline and prohibit multiple calls
+		rctx.m.Lock()
+		alreadyCalled := !rctx.deadline.IsZero()
+		if !alreadyCalled {
+			rctx.deadline = deadline
+		}
+		rctx.m.Unlock()
+		if alreadyCalled {
+			return
+		}
+
+		// Deadline in past?
+		sleepTime := time.Until(deadline)
+		if sleepTime <= 0 {
+			rctx.m.Lock()
+			rctx.err = context.DeadlineExceeded
+			rctx.m.Unlock()
+			close(rctx.done)
+			return
+		}
+		go func() {
+			// Set a timer and wait for timer or parent context to be cancelled
+			timer := time.NewTimer(sleepTime)
+			var setErr error
+			select {
+			case <-pctx.Done():
+				timer.Stop()
+				setErr = pctx.Err()
+			case <-timer.C:
+				setErr = context.DeadlineExceeded
+			}
+			rctx.m.Lock()
+			rctx.err = setErr
+			rctx.m.Unlock()
+			close(rctx.done)
+		}()
+	}
+	return rctx, enforceDeadline
+}

internal/util/optionaldeadline/optionaldeadline_test.go

@@ -0,0 +1,107 @@
+package optionaldeadline
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+
+	"github.com/zrepl/zrepl/internal/util/chainlock"
+	"github.com/zrepl/zrepl/internal/util/zreplcircleci"
+)
+
+func TestContextWithOptionalDeadline(t *testing.T) {
+	zreplcircleci.SkipOnCircleCI(t, "test relies on predictably low scheduler latency")
+
+	ctx := context.Background()
+	cctx, enforceDeadline := ContextWithOptionalDeadline(ctx)
+
+	begin := time.Now()
+	var checker struct {
+		receivedCancellation time.Time
+		cancellationError    error
+		timeout              bool
+		mtx                  chainlock.L
+	}
+	go func() {
+		select {
+		case <-cctx.Done():
+			defer checker.mtx.Lock().Unlock()
+			checker.receivedCancellation = time.Now()
+			checker.cancellationError = cctx.Err()
+		case <-time.After(600 * time.Millisecond):
+			defer checker.mtx.Lock().Unlock()
+			checker.timeout = true
+		}
+	}()
+	defer checker.mtx.Lock().Unlock()
+	checker.mtx.DropWhile(func() {
+		time.Sleep(100 * time.Millisecond)
+	})
+	if !checker.receivedCancellation.IsZero() {
+		t.Fatalf("no enforcement means no cancellation")
+	}
+	require.Nil(t, cctx.Err(), "no error while not cancelled")
+	dl, ok := cctx.Deadline()
+	require.False(t, ok)
+	require.Zero(t, dl)
+	enforceDeadline(begin.Add(200 * time.Millisecond))
+	// second call must be ignored, i.e. we expect the deadline to be at begin+200ms, not begin+400ms
+	enforceDeadline(begin.Add(400 * time.Millisecond))
+
+	checker.mtx.DropWhile(func() {
+		time.Sleep(300 * time.Millisecond) // 100ms margin for scheduler
+	})
+	assert.False(t, checker.timeout, "test timeout")
+	receivedCancellationAfter := checker.receivedCancellation.Sub(begin)
+	if receivedCancellationAfter > 250*time.Millisecond {
+		t.Fatalf("cancellation is beyond acceptable scheduler latency: %s", receivedCancellationAfter)
+	}
+	require.Equal(t, context.DeadlineExceeded, checker.cancellationError)
+}
+
+func TestContextWithOptionalDeadlineNegativeDeadline(t *testing.T) {
+	ctx := context.Background()
+	cctx, enforceDeadline := ContextWithOptionalDeadline(ctx)
+	enforceDeadline(time.Now().Add(-10 * time.Second))
+	select {
+	case <-cctx.Done():
+	default:
+		t.FailNow()
+	}
+}
+
+func TestContextWithOptionalDeadlineParentCancellation(t *testing.T) {
+	zreplcircleci.SkipOnCircleCI(t, "test relies on predictably low scheduler latency")
+
+	pctx, cancel := context.WithCancel(context.Background())
+	cctx, enforceDeadline := ContextWithOptionalDeadline(pctx)
+
+	// 0 ms
+	start := time.Now()
+	enforceDeadline(start.Add(400 * time.Millisecond))
+	time.Sleep(100 * time.Millisecond)
+	cancel()                           // cancel @ ~100ms
+	time.Sleep(100 * time.Millisecond) // give 100ms time to propagate cancel
+	// @ ~200ms
+	select {
+	case <-cctx.Done():
+		assert.True(t, time.Now().Before(start.Add(300*time.Millisecond)))
+		assert.Equal(t, context.Canceled, cctx.Err())
+	default:
+		t.FailNow()
+	}
+
+}
+
+type testContextKey string
+
+const testContextKeyKey testContextKey = "key"
+
+func TestContextWithOptionalDeadlineValue(t *testing.T) {
+	pctx := context.WithValue(context.Background(), testContextKeyKey, "value")
+	cctx, _ := ContextWithOptionalDeadline(pctx)
+	assert.Equal(t, "value", cctx.Value(testContextKeyKey))
+}