Merge pull request #37 from dnozdrin/fix-time-after

Improved memory consumption on timeout checks

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
 	}

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()
 	}
 }

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
 	}
 }