diff --git a/chromedp.go b/chromedp.go index edb80388..2d9b0b39 100644 --- a/chromedp.go +++ b/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 { diff --git a/poll.go b/poll.go index f9307bc1..44cde3e6 100644 --- a/poll.go +++ b/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 { diff --git a/query.go b/query.go index 2239e5ff..8e8e3b18 100644 --- a/query.go +++ b/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.