WIP

util/contextflexibletimeout.go

@@ -0,0 +1,83 @@
+package util
+
+import (
+	"context"
+	"time"
+	"sync"
+)
+
+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 := deadline.Sub(time.Now())
+		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
+}