diff --git a/http/longpoll/longpoll.go b/http/longpoll/longpoll.go index 782b9721c..9b0641b7c 100644 --- a/http/longpoll/longpoll.go +++ b/http/longpoll/longpoll.go @@ -42,51 +42,29 @@ func Until(ctx context.Context, d time.Duration, fn LongPollFunc) (ok bool, err // budget is communicated via the provided context. This is a defence measure to not have accidental // long running routines. If no duration is given (0) the long poll will have exactly one execution. func (c Config) LongPollUntil(ctx context.Context, d time.Duration, fn LongPollFunc) (ok bool, err error) { - until := time.Now() + var cancel context.CancelFunc if d != 0 { - if d < c.MinWaitTime { // guard lower bound - until = until.Add(c.MinWaitTime) - } else if d > c.MaxWaitTime { // guard upper bound - until = until.Add(c.MaxWaitTime) - } else { - until = until.Add(d) + if d < c.MinWaitTime { + d = c.MinWaitTime + } else if d > c.MaxWaitTime { + d = c.MaxWaitTime } - } - fnCtx, cancel := context.WithDeadline(ctx, until) - defer cancel() + ctx, cancel = context.WithTimeout(ctx, d) + defer cancel() + } -loop: for { - ok, err = fn(fnCtx) - if err != nil { + ok, err = fn(ctx) + if ok || err != nil || d == 0 { return } - // fn returns true, break the loop - if ok { - break - } - - // no long polling - if d <= 0 { - break - } - - // long pooling desired? - if !time.Now().Add(c.RetryTime).Before(until) { - break - } - select { - // handle context cancelation case <-ctx.Done(): return false, ctx.Err() case <-time.After(c.RetryTime): - continue loop } } - - return } diff --git a/http/longpoll/longpoll_test.go b/http/longpoll/longpoll_test.go index 289860ab9..a0c80534e 100644 --- a/http/longpoll/longpoll_test.go +++ b/http/longpoll/longpoll_test.go @@ -37,7 +37,13 @@ func TestLongPollUntilBounds(t *testing.T) { func TestLongPollUntilNoTimeout(t *testing.T) { called := 0 - ok, err := Until(context.Background(), 0, func(context.Context) (bool, error) { + ok, err := Until(context.Background(), 0, func(ctx context.Context) (bool, error) { + f := func(ctx context.Context) { + assert.NoError(t, ctx.Err()) + } + + f(ctx) + called++ return false, nil }) @@ -73,8 +79,8 @@ func TestLongPollUntilTimeout(t *testing.T) { return false, nil }) assert.False(t, ok) - assert.NoError(t, err) - assert.Equal(t, 2, called) + assert.Error(t, err) + assert.GreaterOrEqual(t, called, 2) } func TestLongPollUntilTimeoutWithContext(t *testing.T) {