context.go

   1 package signalcontext
   2
   3 import (
   4         "context"
   5         "os"
   6         "os/signal"
   7         "sync"
   8         "time"
   9 )
  10
  11 // A Context is an implementation of the context.Context interface which
  12 // cancels when an operating system signal (e.g. os.Interrupt) is
  13 // received.
  14 //
  15 // Contexts should be created via the UntilSignal function.
  16 type Context struct {
  17         parent context.Context
  18         done   chan struct{}
  19         err    error
  20
  21         // The mutex synchronizes access to err and clearing the
  22         // internal Signal channel after initialization.
  23         m sync.Mutex
  24         c chan os.Signal
  25 }
  26
  27 // UntilSignal returns a new Context which will cancel when the parent
  28 // does or when any of the specified operating system signals are
  29 // received.
  30 func UntilSignal(parent context.Context, sig ...os.Signal) *Context {
  31         ctx := new(Context)
  32         ctx.parent = parent
  33
  34         if err := parent.Err(); err != nil {
  35                 ctx.done = alreadyclosed
  36                 ctx.err = err
  37                 return ctx
  38         }
  39
  40         ctx.done = make(chan struct{})
  41         ctx.c = make(chan os.Signal, 1)
  42         signal.Notify(ctx.c, sig...)
  43         go ctx.wait(sig...)
  44         return ctx
  45 }
  46
  47 func (s *Context) wait(sig ...os.Signal) {
  48         var err error
  49         select {
  50         case <-s.parent.Done():
  51                 err = s.parent.Err()
  52         case v := <-s.c:
  53                 if v != nil {
  54                         err = Error{v}
  55                 }
  56         }
  57         signal.Stop(s.c)
  58         s.m.Lock()
  59         if s.err == nil {
  60                 s.err = err
  61         }
  62         close(s.c)
  63         s.c = nil
  64         s.m.Unlock()
  65         close(s.done)
  66 }
  67
  68 // Cancel cancels this context, if it hasn’t already been. (If it has,
  69 // this is safe but has no effect.) Canceling this context releases
  70 // resources associated with it and stops listening for the configured
  71 // operating system signals.
  72 func (s *Context) Cancel() {
  73         s.m.Lock()
  74         if s.c != nil {
  75                 s.err = context.Canceled
  76                 select {
  77                 case s.c <- nil:
  78                 default:
  79                 }
  80         }
  81         s.m.Unlock()
  82 }
  83
  84 // Deadline implements context.Context; a Context’s deadline is that of
  85 // its parent.
  86 func (s *Context) Deadline() (time.Time, bool) {
  87         return s.parent.Deadline()
  88 }
  89
  90 // Value implements context.Context; any Context value is that of its
  91 // parent.
  92 func (s *Context) Value(key interface{}) interface{} {
  93         return s.parent.Value(key)
  94 }
  95
  96 // Done returns a channel that's closed when work done on behalf of this
  97 // context should be canceled, either because the parent finished or
  98 // because one of the configured operating system signals was received.
  99 func (s *Context) Done() <-chan struct{} {
 100         return s.done
 101 }
 102
 103 // Err implements context.Context; it returns context.Canceled if the
 104 // context was canceled by its Cancel method; an Error if the context
 105 // canceled due to a operating system signal; the parent’s error if the
 106 // parent canceled before either of those; or nil if the context is not
 107 // yet canceled.
 108 func (s *Context) Err() error {
 109         s.m.Lock()
 110         err := s.err
 111         s.m.Unlock()
 112         return err
 113 }
 114
 115 // Reuse the same channel for all Contexts which begin life already
 116 // canceled.
 117 var alreadyclosed = make(chan struct{})
 118
 119 func init() {
 120         close(alreadyclosed)
 121 }