-
Notifications
You must be signed in to change notification settings - Fork 60
Commit
Signed-off-by: Joe Lanford <[email protected]>
- Loading branch information
There are no files selected for viewing
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,113 @@ | ||
package context | ||
|
||
import ( | ||
"context" | ||
"time" | ||
) | ||
|
||
func (d *delayContext) Deadline() (time.Time, bool) { | ||
select { | ||
case <-d.parentCtx.Done(): | ||
// if the parent context is done, wait | ||
// for our timeout setup to complete, then | ||
// return the timeout context's deadline. | ||
<-d.setupDone | ||
return d.timeoutCtx.Deadline() | ||
default: | ||
// if the parent context has a deadline, simply add | ||
// our delay. | ||
if parentDeadline, ok := d.parentCtx.Deadline(); ok { | ||
return parentDeadline.Add(d.delay), true | ||
} | ||
// if the parent context does not have a deadline | ||
// then we don't know ours either because it depends | ||
// on when the parent is done. | ||
return time.Time{}, false | ||
} | ||
} | ||
|
||
func (d *delayContext) Done() <-chan struct{} { | ||
return d.done | ||
} | ||
|
||
func (d *delayContext) Err() error { | ||
// If the parent context is done, wait until setup | ||
// is done, then return the timeout context's error. | ||
select { | ||
case <-d.parentCtx.Done(): | ||
<-d.setupDone | ||
return d.timeoutCtx.Err() | ||
default: | ||
} | ||
|
||
// If done is closed, that means we were | ||
// directly cancelled. Otherwise (if neither | ||
// parent context is done or done is closed) | ||
// the context is still active, hence no error | ||
select { | ||
case <-d.done: | ||
return context.Canceled | ||
default: | ||
return nil | ||
} | ||
} | ||
|
||
func (d *delayContext) Value(key interface{}) interface{} { | ||
return d.parentCtx.Value(key) | ||
} | ||
|
||
type delayContext struct { | ||
parentCtx context.Context | ||
delay time.Duration | ||
|
||
done chan struct{} | ||
setupDone chan struct{} | ||
|
||
timeoutCtx context.Context | ||
timeoutCancel context.CancelFunc | ||
} | ||
|
||
func WithDelay(parentCtx context.Context, delay time.Duration) (context.Context, context.CancelFunc) { | ||
delayedCtx := &delayContext{ | ||
parentCtx: parentCtx, | ||
delay: delay, | ||
done: make(chan struct{}), | ||
setupDone: make(chan struct{}), | ||
} | ||
|
||
setupDelay := func() { | ||
timeoutCtx, timeoutCancel := context.WithTimeout(context.Background(), delay) | ||
context.AfterFunc(timeoutCtx, func() { close(delayedCtx.done) }) | ||
delayedCtx.timeoutCtx = timeoutCtx | ||
delayedCtx.timeoutCancel = timeoutCancel | ||
close(delayedCtx.setupDone) | ||
} | ||
|
||
unregisterDelay := context.AfterFunc(parentCtx, setupDelay) | ||
|
||
cancelFunc := func() { | ||
setupNeverHappened := unregisterDelay() | ||
if setupNeverHappened { | ||
// if setup never happened, then the delay context was | ||
// cancelled prior to the parent context being done. | ||
// | ||
// all we need to do here is close the done chan. | ||
close(delayedCtx.done) | ||
} else { | ||
// if we're here, the setup function was called | ||
|
||
// wait until setup is done to ensure there is a | ||
// timeoutContext/timeoutCancel | ||
<-delayedCtx.setupDone | ||
|
||
// cancel the timeout context (which includes | ||
// an AfterFunc to also close our doneChan, so | ||
// we'll wait for that to be closed before | ||
// returning) | ||
delayedCtx.timeoutCancel() | ||
<-delayedCtx.done | ||
} | ||
} | ||
|
||
return delayedCtx, cancelFunc | ||
} |
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,99 @@ | ||
package context_test | ||
|
||
import ( | ||
"context" | ||
"testing" | ||
"time" | ||
|
||
"github.com/stretchr/testify/assert" | ||
|
||
contextutil "github.com/operator-framework/operator-controller/internal/util/context" | ||
) | ||
|
||
func TestWithDelay_Delays(t *testing.T) { | ||
for _, delay := range []time.Duration{ | ||
0, | ||
time.Millisecond * 10, | ||
time.Millisecond * 100, | ||
time.Millisecond * 200, | ||
} { | ||
t.Run(delay.String(), func(t *testing.T) { | ||
parentCtx, parentCancel := context.WithCancel(context.Background()) | ||
delayCtx, _ := contextutil.WithDelay(parentCtx, delay) | ||
|
||
parentCancel() | ||
|
||
// verify deadline is within 1m ms of what we expect | ||
expectDeadline := time.Now().Add(delay) | ||
actualDeadline, ok := delayCtx.Deadline() | ||
assert.True(t, ok, "expected delay context to have a deadline after parent was cancelled") | ||
assert.WithinDurationf(t, expectDeadline, actualDeadline, time.Millisecond, "expected the context's deadline (%v) to be within 1 ms of %v; diff was %v", expectDeadline, actualDeadline, expectDeadline.Sub(actualDeadline)) | ||
|
||
// verify context is done due to deadline exceeded and that it happens | ||
// within 3ms of our expectation | ||
select { | ||
case <-delayCtx.Done(): | ||
assert.ErrorIs(t, delayCtx.Err(), context.DeadlineExceeded) | ||
case <-time.After(expectDeadline.Sub(time.Now()) + 3*time.Millisecond): | ||
diff := time.Now().Sub(expectDeadline) | ||
t.Fatalf("delay context should have been canceled quickly after %s, but it took %s", delay, diff) | ||
} | ||
}) | ||
} | ||
} | ||
|
||
func TestWithDelay_Deadline(t *testing.T) { | ||
t.Run("parent has deadline", func(t *testing.T) { | ||
parentDeadline := time.Now().Add(200 * time.Millisecond) | ||
parentCtx, _ := context.WithDeadline(context.Background(), parentDeadline) | ||
Check failure on line 48 in internal/util/context/context_test.go
|
||
|
||
delay := 250 * time.Millisecond | ||
delayCtx, _ := contextutil.WithDelay(parentCtx, delay) | ||
|
||
expectDeadline := parentDeadline.Add(delay) | ||
actualDeadline, ok := delayCtx.Deadline() | ||
|
||
assert.True(t, ok, "expected delay context to have a deadline before parent was cancelled") | ||
assert.Equal(t, expectDeadline, actualDeadline) | ||
}) | ||
t.Run("parent has no deadline", func(t *testing.T) { | ||
parentCtx, _ := context.WithCancel(context.Background()) | ||
Check failure on line 60 in internal/util/context/context_test.go
|
||
|
||
delayCtx, _ := contextutil.WithDelay(parentCtx, 200*time.Millisecond) | ||
actualDeadline, ok := delayCtx.Deadline() | ||
assert.False(t, ok, "expected delay context to have an unknown deadline before parent was cancelled") | ||
assert.Equal(t, time.Time{}, actualDeadline, "expected delay context deadline to be unset") | ||
|
||
}) | ||
} | ||
|
||
func TestWithDelay_Err(t *testing.T) { | ||
t.Run("nil", func(t *testing.T) { | ||
delayCtx, _ := contextutil.WithDelay(context.Background(), 0) | ||
assert.NoError(t, delayCtx.Err()) | ||
}) | ||
t.Run("canceled before parent done", func(t *testing.T) { | ||
delayCtx, delayCancel := contextutil.WithDelay(context.Background(), 0) | ||
delayCancel() | ||
assert.ErrorIs(t, delayCtx.Err(), context.Canceled) | ||
}) | ||
t.Run("canceled after parent done", func(t *testing.T) { | ||
parentCtx, parentCancel := context.WithCancel(context.Background()) | ||
delayCtx, delayCancel := contextutil.WithDelay(parentCtx, 200*time.Millisecond) | ||
parentCancel() | ||
delayCancel() | ||
assert.ErrorIs(t, delayCtx.Err(), context.Canceled) | ||
}) | ||
t.Run("deadline exceeded", func(t *testing.T) { | ||
parentCtx, parentCancel := context.WithCancel(context.Background()) | ||
delayCtx, _ := contextutil.WithDelay(parentCtx, 0) | ||
parentCancel() | ||
assert.ErrorIs(t, delayCtx.Err(), context.DeadlineExceeded) | ||
}) | ||
} | ||
|
||
func TestWithDelay_Value(t *testing.T) { | ||
parentCtx := context.WithValue(context.Background(), "foo", "bar") | ||
delayCtx, _ := contextutil.WithDelay(parentCtx, 0) | ||
assert.Equal(t, "bar", delayCtx.Value("foo")) | ||
} |