package signalcontext import ( "context" "os" "os/signal" "sync" "time" ) // A Context is an implementation of the context.Context interface which // cancels when an operating system signal (e.g. os.Interrupt) is // received. // // Contexts should be created via the UntilSignal function. type Context struct { parent context.Context done chan struct{} err error // The mutex synchronizes access to err and clearing the // internal Signal channel after initialization. m sync.Mutex c chan os.Signal } // UntilSignal returns a new Context which will cancel when the parent // does or when any of the specified operating system signals are // received. func UntilSignal(parent context.Context, sig ...os.Signal) *Context { ctx := new(Context) ctx.parent = parent if err := parent.Err(); err != nil { ctx.done = alreadyclosed ctx.err = err return ctx } ctx.done = make(chan struct{}) ctx.c = make(chan os.Signal, 1) signal.Notify(ctx.c, sig...) go ctx.wait(sig...) return ctx } func (s *Context) wait(sig ...os.Signal) { var err error select { case <-s.parent.Done(): err = s.parent.Err() case v := <-s.c: if v != nil { err = Error{v} } } signal.Stop(s.c) s.m.Lock() if s.err == nil { s.err = err } close(s.c) s.c = nil s.m.Unlock() close(s.done) } // Cancel cancels this context, if it hasn’t already been. (If it has, // this is safe but has no effect.) Canceling this context releases // resources associated with it and stops listening for the configured // operating system signals. func (s *Context) Cancel() { s.m.Lock() if s.c != nil { s.err = context.Canceled select { case s.c <- nil: default: } } s.m.Unlock() } // Deadline implements context.Context; a Context’s deadline is that of // its parent. func (s *Context) Deadline() (time.Time, bool) { return s.parent.Deadline() } // Value implements context.Context; any Context value is that of its // parent. func (s *Context) Value(key interface{}) interface{} { return s.parent.Value(key) } // Done returns a channel that's closed when work done on behalf of this // context should be canceled, either because the parent finished or // because one of the configured operating system signals was received. func (s *Context) Done() <-chan struct{} { return s.done } // Err implements context.Context; it returns context.Canceled if the // context was canceled by its Cancel method; an Error if the context // canceled due to a operating system signal; the parent’s error if the // parent canceled before either of those; or nil if the context is not // yet canceled. func (s *Context) Err() error { s.m.Lock() err := s.err s.m.Unlock() return err } // Reuse the same channel for all Contexts which begin life already // canceled. var alreadyclosed = make(chan struct{}) func init() { close(alreadyclosed) }