http/longpoll: Fix bug and simplify implementation

This fixes a bug where if duration was set to 0 the context would expire
instantly causing any context-aware function to stop.

Author: monstermunchkin

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
 }

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) {