github.com/mailgun/holster/v4@v4.20.0/cancel/context.go (about)

     1  package cancel
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  	"time"
     7  )
     8  
     9  type Context interface {
    10  	context.Context
    11  	Wrap(context.Context) context.Context
    12  	Cancel()
    13  }
    14  
    15  type cancelCtx struct {
    16  	ctx    context.Context
    17  	cancel context.CancelFunc
    18  }
    19  
    20  // New creates a context that wraps the given context and returns an obj that can be cancelled.
    21  // This allows an object which desires to cancel a long running operation to store a single
    22  // cancel.Context in it's struct variables instead of having to store both the context.Context
    23  // and context.CancelFunc.
    24  func New(ctx context.Context) Context {
    25  	if ctx == nil {
    26  		ctx = context.Background()
    27  	}
    28  
    29  	ctx, cancel := context.WithCancel(ctx)
    30  	return &cancelCtx{
    31  		cancel: cancel,
    32  		ctx:    ctx,
    33  	}
    34  }
    35  
    36  func (c *cancelCtx) Cancel()                                 { c.cancel() }
    37  func (c *cancelCtx) Deadline() (deadline time.Time, ok bool) { return c.ctx.Deadline() }
    38  func (c *cancelCtx) Done() <-chan struct{}                   { return c.ctx.Done() }
    39  func (c *cancelCtx) Err() error                              { return c.ctx.Err() }
    40  func (c *cancelCtx) Value(key interface{}) interface{}       { return c.ctx.Value(key) }
    41  
    42  // Wrap returns a Context that will be cancelled when either cancel.Context or the passed context is cancelled
    43  func (c *cancelCtx) Wrap(ctx context.Context) context.Context {
    44  	return NewWrappedContext(ctx, c)
    45  }
    46  
    47  // NewWrappedContext returns a Context that will be cancelled when either of the passed contexts are cancelled
    48  func NewWrappedContext(left, right context.Context) context.Context {
    49  	w := WrappedContext{
    50  		doneCh: make(chan struct{}),
    51  	}
    52  	// Wait for either ctx to be cancelled and propagate to the wrapped context
    53  	go func() {
    54  		select {
    55  		case <-left.Done():
    56  			w.Reason(left.Err())
    57  		case <-right.Done():
    58  			w.Reason(right.Err())
    59  		}
    60  	}()
    61  	return &w
    62  }
    63  
    64  type WrappedContext struct {
    65  	context.Context
    66  
    67  	mutex  sync.Mutex
    68  	doneCh chan struct{}
    69  	err    error
    70  }
    71  
    72  func (w *WrappedContext) Reason(err error) {
    73  	w.mutex.Lock()
    74  	defer w.mutex.Unlock()
    75  	close(w.doneCh)
    76  	w.err = err
    77  }
    78  
    79  func (w *WrappedContext) Done() <-chan struct{} {
    80  	w.mutex.Lock()
    81  	defer w.mutex.Unlock()
    82  	return w.doneCh
    83  }
    84  
    85  func (w *WrappedContext) Err() error {
    86  	w.mutex.Lock()
    87  	defer w.mutex.Unlock()
    88  	return w.err
    89  }