diff --git a/deadline/deadline.go b/deadline/deadline.go index c97fd86..cff6430 100644 --- a/deadline/deadline.go +++ b/deadline/deadline.go @@ -39,10 +39,15 @@ func (d *Deadline) Run(work func(<-chan struct{}) error) error { } }() + timer := time.NewTimer(d.timeout) select { case ret := <-result: + if !timer.Stop() { + <-timer.C + } + return ret - case <-time.After(d.timeout): + case <-timer.C: close(stopper) return ErrTimedOut } diff --git a/retrier/retrier.go b/retrier/retrier.go index 6801c55..1cd8d47 100644 --- a/retrier/retrier.go +++ b/retrier/retrier.go @@ -60,8 +60,8 @@ func (r *Retrier) RunCtx(ctx context.Context, work func(ctx context.Context) err return ret } - timeout := time.After(r.calcSleep(retries)) - if err := r.sleep(ctx, timeout); err != nil { + timer := time.NewTimer(r.calcSleep(retries)) + if err := r.sleep(ctx, timer); err != nil { return err } @@ -70,11 +70,15 @@ func (r *Retrier) RunCtx(ctx context.Context, work func(ctx context.Context) err } } -func (r *Retrier) sleep(ctx context.Context, t <-chan time.Time) error { +func (r *Retrier) sleep(ctx context.Context, timer *time.Timer) error { select { - case <-t: + case <-timer.C: return nil case <-ctx.Done(): + if !timer.Stop() { + <-timer.C + } + return ctx.Err() } } diff --git a/semaphore/semaphore.go b/semaphore/semaphore.go index d4aaa04..b3c1b93 100644 --- a/semaphore/semaphore.go +++ b/semaphore/semaphore.go @@ -29,10 +29,15 @@ func New(tickets int, timeout time.Duration) *Semaphore { // If it cannot after "timeout" amount of time, it returns ErrNoTickets. It is // safe to call Acquire concurrently on a single Semaphore. func (s *Semaphore) Acquire() error { + timer := time.NewTimer(s.timeout) select { case s.sem <- struct{}{}: + if !timer.Stop() { + <-timer.C + } + return nil - case <-time.After(s.timeout): + case <-timer.C: return ErrNoTickets } }