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,43 @@ 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
 		}
+
+		ctx, cancel = context.WithTimeout(ctx, d)
+		defer cancel()
 	}
 
-	fnCtx, cancel := context.WithDeadline(ctx, until)
-	defer cancel()
+	doneCh := make(chan struct{})
 
-loop:
-	for {
-		ok, err = fn(fnCtx)
-		if err != nil {
-			return
-		}
+	go func() {
+		defer func() {
+			doneCh <- struct{}{}
+		}()
 
-		// fn returns true, break the loop
-		if ok {
-			break
-		}
+		for {
+			ok, err = fn(ctx)
+			if ok || err != nil || d == 0 {
+				return
+			}
 
-		// no long polling
-		if d <= 0 {
-			break
-		}
+			select {
+			case <-ctx.Done():
+				err = ctx.Err()
 
-		// long pooling desired?
-		if !time.Now().Add(c.RetryTime).Before(until) {
-			break
+				return
+			case <-time.After(c.RetryTime):
+			}
 		}
+	}()
 
-		select {
-		// handle context cancelation
-		case <-ctx.Done():
-			return false, ctx.Err()
-		case <-time.After(c.RetryTime):
-			continue loop
-		}
-	}
+	<-doneCh
 
 	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) {