avoid the unnecessary wait in the first iteration of a loop

fixes #928

chromedp.go

@@ -627,17 +627,38 @@ func (t Tasks) Do(ctx context.Context) error {
 // been able to be written/tested.
 func Sleep(d time.Duration) Action {
 	return ActionFunc(func(ctx context.Context) error {
-		// Don't use time.After, to avoid a temporary goroutine leak if
-		// ctx is cancelled before the timer fires.
-		t := time.NewTimer(d)
-		select {
-		case <-ctx.Done():
-			t.Stop()
-			return ctx.Err()
-		case <-t.C:
+		return sleepContext(ctx, d)
+	})
+}
+
+// sleepContext sleeps for the specified duration. It returns ctx.Err() immediately
+// if the context is cancelled.
+func sleepContext(ctx context.Context, d time.Duration) error {
+	timer := time.NewTimer(d)
+	select {
+	case <-ctx.Done():
+		if !timer.Stop() {
+			<-timer.C
 		}
+		return ctx.Err()
+	case <-timer.C:
 		return nil
-	})
+	}
+}
+
+// retryWithSleep reties the execution of the specified func until the func returns
+// true (means to stop) or a non-nil error.
+func retryWithSleep(ctx context.Context, d time.Duration, f func(ctx context.Context) (bool, error)) error {
+	for {
+		toStop, err := f(ctx)
+		if toStop || err != nil {
+			return err
+		}
+		err = sleepContext(ctx, d)
+		if err != nil {
+			return err
+		}
+	}
 }
 
 type cancelableListener struct {

poll.go

@@ -38,13 +38,15 @@ func (p *pollTask) Do(ctx context.Context) error {
 		execCtx runtime.ExecutionContextID
 		ok      bool
 	)
-	for !ok {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-time.After(5 * time.Millisecond):
-		}
+
+	for {
 		_, _, execCtx, ok = t.ensureFrame()
+		if ok {
+			break
+		}
+		if err := sleepContext(ctx, 5*time.Millisecond); err != nil {
+			return err
+		}
 	}
 
 	if p.frame != nil {

query.go

@@ -155,16 +155,10 @@ func (s *Selector) Do(ctx context.Context) error {
 	if t == nil {
 		return ErrInvalidTarget
 	}
-	for {
-		select {
-		case <-ctx.Done():
-			return ctx.Err()
-		case <-time.After(5 * time.Millisecond):
-		}
-
+	return retryWithSleep(ctx, 5*time.Millisecond, func(ctx context.Context) (bool, error) {
 		frame, root, execCtx, ok := t.ensureFrame()
 		if !ok {
-			continue
+			return false, nil
 		}
 
 		fromNode := s.fromNode
@@ -194,20 +188,20 @@ func (s *Selector) Do(ctx context.Context) error {
 
 		ids, err := s.by(ctx, fromNode)
 		if err != nil || len(ids) < s.exp {
-			continue
+			return false, nil
 		}
 		nodes, err := s.wait(ctx, frame, execCtx, ids...)
 		// if nodes==nil, we're not yet ready
 		if nodes == nil || err != nil {
-			continue
+			return false, nil
 		}
 		if s.after != nil {
 			if err := s.after(ctx, execCtx, nodes...); err != nil {
-				return err
+				return true, err
 			}
 		}
-		return nil
-	}
+		return true, nil
+	})
 }
 
 // selAsString forces sel into a string.