context.go raw

   1  package clock
   2  
   3  import (
   4  	"context"
   5  	"fmt"
   6  	"sync"
   7  	"time"
   8  )
   9  
  10  func (m *Mock) WithTimeout(parent context.Context, timeout time.Duration) (context.Context, context.CancelFunc) {
  11  	return m.WithDeadline(parent, m.Now().Add(timeout))
  12  }
  13  
  14  func (m *Mock) WithDeadline(parent context.Context, deadline time.Time) (context.Context, context.CancelFunc) {
  15  	if cur, ok := parent.Deadline(); ok && cur.Before(deadline) {
  16  		// The current deadline is already sooner than the new one.
  17  		return context.WithCancel(parent)
  18  	}
  19  	ctx := &timerCtx{clock: m, parent: parent, deadline: deadline, done: make(chan struct{})}
  20  	propagateCancel(parent, ctx)
  21  	dur := m.Until(deadline)
  22  	if dur <= 0 {
  23  		ctx.cancel(context.DeadlineExceeded) // deadline has already passed
  24  		return ctx, func() {}
  25  	}
  26  	ctx.Lock()
  27  	defer ctx.Unlock()
  28  	if ctx.err == nil {
  29  		ctx.timer = m.AfterFunc(dur, func() {
  30  			ctx.cancel(context.DeadlineExceeded)
  31  		})
  32  	}
  33  	return ctx, func() { ctx.cancel(context.Canceled) }
  34  }
  35  
  36  // propagateCancel arranges for child to be canceled when parent is.
  37  func propagateCancel(parent context.Context, child *timerCtx) {
  38  	if parent.Done() == nil {
  39  		return // parent is never canceled
  40  	}
  41  	go func() {
  42  		select {
  43  		case <-parent.Done():
  44  			child.cancel(parent.Err())
  45  		case <-child.Done():
  46  		}
  47  	}()
  48  }
  49  
  50  type timerCtx struct {
  51  	sync.Mutex
  52  
  53  	clock    Clock
  54  	parent   context.Context
  55  	deadline time.Time
  56  	done     chan struct{}
  57  
  58  	err   error
  59  	timer *Timer
  60  }
  61  
  62  func (c *timerCtx) cancel(err error) {
  63  	c.Lock()
  64  	defer c.Unlock()
  65  	if c.err != nil {
  66  		return // already canceled
  67  	}
  68  	c.err = err
  69  	close(c.done)
  70  	if c.timer != nil {
  71  		c.timer.Stop()
  72  		c.timer = nil
  73  	}
  74  }
  75  
  76  func (c *timerCtx) Deadline() (deadline time.Time, ok bool) { return c.deadline, true }
  77  
  78  func (c *timerCtx) Done() <-chan struct{} { return c.done }
  79  
  80  func (c *timerCtx) Err() error { return c.err }
  81  
  82  func (c *timerCtx) Value(key interface{}) interface{} { return c.parent.Value(key) }
  83  
  84  func (c *timerCtx) String() string {
  85  	return fmt.Sprintf("clock.WithDeadline(%s [%s])", c.deadline, c.deadline.Sub(c.clock.Now()))
  86  }
  87