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 // completes when a signal (e.g. os.Interrupt) is received.
  13 //
  14 // Contexts should be created via the UntilSignal function.
  15 type Context struct {
  16         parent context.Context
  17         done   chan struct{}
  18         err    error
  19
  20         // The mutex synchronizes access to err and clearing the
  21         // internal Signal channel after initialization.
  22         m sync.Mutex
  23         c chan os.Signal
  24 }
  25
  26 // UntilSignal returns a new Context which will complete when the parent
  27 // does or when any of the specified signals are received.
  28 func UntilSignal(parent context.Context, sig ...os.Signal) *Context {
  29         ctx := new(Context)
  30         ctx.parent = parent
  31         ctx.done = make(chan struct{})
  32
  33         if err := parent.Err(); err != nil {
  34                 close(ctx.done)
  35                 ctx.err = err
  36                 return ctx
  37         }
  38
  39         ctx.c = make(chan os.Signal, 1)
  40         signal.Notify(ctx.c, sig...)
  41         go ctx.wait(sig...)
  42         return ctx
  43 }
  44
  45 func (s *Context) wait(sig ...os.Signal) {
  46         var err error
  47         select {
  48         case <-s.parent.Done():
  49                 err = s.parent.Err()
  50         case v := <-s.c:
  51                 if v != nil {
  52                         err = Error{v}
  53                 }
  54         }
  55         signal.Stop(s.c)
  56         s.m.Lock()
  57         if s.err == nil {
  58                 s.err = err
  59         }
  60         close(s.c)
  61         s.c = nil
  62         s.m.Unlock()
  63         close(s.done)
  64 }
  65
  66 // Cancel cancels this context, if it hasn’t already been completed. (If
  67 // it has, this is safe but has no effect.)
  68 func (s *Context) Cancel() {
  69         s.m.Lock()
  70         if s.c != nil {
  71                 s.err = context.Canceled
  72                 select {
  73                 case s.c <- nil:
  74                 default:
  75                 }
  76         }
  77         s.m.Unlock()
  78 }
  79
  80 // Deadline implements context.Context; a Context’s deadline is that of
  81 // its parent.
  82 func (s *Context) Deadline() (time.Time, bool) {
  83         return s.parent.Deadline()
  84 }
  85
  86 // Value implements context.Context; any value is that of its parent.
  87 func (s *Context) Value(key interface{}) interface{} {
  88         return s.parent.Value(key)
  89 }
  90
  91 // Done implements context.Context.
  92 func (s *Context) Done() <-chan struct{} {
  93         return s.done
  94 }
  95
  96 // Err implements context.Context; it returns context.Canceled if the
  97 // context was canceled; an Error if the context completed due to a
  98 // signal; the parent’s error if the parent was done before either of
  99 // those; or nil if the context is not yet done.
 100 func (s *Context) Err() error {
 101         s.m.Lock()
 102         err := s.err
 103         s.m.Unlock()
 104         return err
 105 }