From bdcea55117ecc04722b27d74f2d5a2203fffae6d Mon Sep 17 00:00:00 2001 From: Mitchell Hashimoto Date: Mon, 15 Aug 2016 20:54:25 -0700 Subject: [PATCH] helper/signalwrapper --- helper/signalwrapper/wrapper.go | 177 +++++++++++++++++++++++++++ helper/signalwrapper/wrapper_test.go | 68 ++++++++++ 2 files changed, 245 insertions(+) create mode 100644 helper/signalwrapper/wrapper.go create mode 100644 helper/signalwrapper/wrapper_test.go diff --git a/helper/signalwrapper/wrapper.go b/helper/signalwrapper/wrapper.go new file mode 100644 index 000000000..765661ece --- /dev/null +++ b/helper/signalwrapper/wrapper.go @@ -0,0 +1,177 @@ +// Package signalwrapper is used to run functions that are sensitive to +// signals that may be received from outside the process. It can also be +// used as just an async function runner that is cancellable and we may +// abstract this further into another package in the future. +package signalwrapper + +import ( + "os" + "os/signal" + "sync" +) + +// CancellableFunc is a function that cancels if it receives a message +// on the given channel. It must return an error if any occurred. It should +// return no error if it was cancelled successfully since it is assumed +// that this function will probably be called again at some future point +// since it was interrupted. +type CancellableFunc func(<-chan struct{}) error + +// Run wraps and runs the given cancellable function and returns the Wrapped +// struct that can be used to listen for events, cancel on other events +// (such as timeouts), etc. +func Run(f CancellableFunc) *Wrapped { + // Determine the signals we're listening to. Prematurely making + // this a slice since I predict a future where we'll add others and + // the complexity in doing so is low. + signals := []os.Signal{os.Interrupt} + + // Register a listener for the signals + sigCh := make(chan os.Signal, 1) + signal.Notify(sigCh, signals...) + + // Create the channel we'll use to "cancel" + cancelCh := make(chan struct{}) + + // This channel keeps track of whether the function we're running + // completed successfully and the errors it may have had. It is + // VERY IMPORTANT that the errCh is buffered to at least 1 so that + // it doesn't block when finishing. + doneCh := make(chan struct{}) + errCh := make(chan error, 1) + + // Build our wrapped result + wrapped := &Wrapped{ + ErrCh: errCh, + errCh: errCh, + cancelCh: cancelCh, + } + + // Start the function + go func() { + err := f(cancelCh) + + // Close the done channel _before_ sending the error in case + // the error channel read blocks (it shouldn't) to avoid interrupts + // doing anything. + close(doneCh) + + // Mark completion + wrapped.done(err) + }() + + // Goroutine to track interrupts and make sure we do at-most-once + // delivery of an interrupt since we're using a channel. + go func() { + // Clean up after this since this is the only function that + // reads signals. + defer signal.Stop(sigCh) + + select { + case <-doneCh: + // Everything happened naturally + case <-sigCh: + // Stop the function. Goroutine since we don't care about + // the result and we'd like to end this goroutine as soon + // as possible to avoid any more signals coming in. + go wrapped.Cancel() + } + }() + + return wrapped +} + +// Wrapped is the return value of wrapping a function. This has channels +// that can be used to wait for a result as well as functions to help with +// different behaviors. +type Wrapped struct { + // Set and consumed by user + + // ErrCh is the channel to listen for real-time events on the wrapped + // function. A nil error sent means the execution completed without error. + // This is an exactly once delivery channel. + ErrCh <-chan error + + // Set by creator + errCh chan<- error + cancelCh chan<- struct{} + + // Set automatically + once sync.Once + cancelCond *sync.Cond + cancelLock *sync.Mutex + resultErr error + resultSet bool +} + +// Cancel stops the running function and blocks until it returns. The +// resulting value is returned. +// +// It is safe to call this multiple times. This will return the resulting +// error value each time. +func (w *Wrapped) Cancel() error { + w.once.Do(w.init) + w.cancelLock.Lock() + + // If we have a result set, return that + if w.resultSet { + w.cancelLock.Unlock() + return w.resultErr + } + + // If we have a cancel channel, close it to signal and set it to + // nil so we never do that again. + if w.cancelCh != nil { + close(w.cancelCh) + w.cancelCh = nil + } + + // Wait for the result to be set + defer w.cancelLock.Unlock() + w.cancelCond.Wait() + return w.resultErr +} + +// Wait waits for the completion of the wrapped function and returns the result. +// +// This can be called multiple times safely. +func (w *Wrapped) Wait() error { + w.once.Do(w.init) + w.cancelLock.Lock() + defer w.cancelLock.Unlock() + + // If we don't have a result yet, wait for that + if !w.resultSet { + w.cancelCond.Wait() + } + + // Return the result + return w.resultErr +} + +// done marks this wrapped function as done with the resulting value. +// This must only be called once. +func (w *Wrapped) done(err error) { + w.once.Do(w.init) + w.cancelLock.Lock() + + // Set the result + w.resultErr = err + w.resultSet = true + + // Notify any waiters + w.cancelCond.Broadcast() + + // Unlock since the next call can be blocking + w.cancelLock.Unlock() + + // Notify any channel listeners + w.errCh <- err +} + +func (w *Wrapped) init() { + // Create the condition variable + var m sync.Mutex + w.cancelCond = sync.NewCond(&m) + w.cancelLock = &m +} diff --git a/helper/signalwrapper/wrapper_test.go b/helper/signalwrapper/wrapper_test.go new file mode 100644 index 000000000..4efa1a146 --- /dev/null +++ b/helper/signalwrapper/wrapper_test.go @@ -0,0 +1,68 @@ +package signalwrapper + +import ( + "errors" + "testing" + "time" +) + +func TestWrapped_goodCh(t *testing.T) { + errVal := errors.New("hi") + f := func(<-chan struct{}) error { return errVal } + err := <-Run(f).ErrCh + if err != errVal { + t.Fatalf("bad: %#v", err) + } +} + +func TestWrapped_goodWait(t *testing.T) { + errVal := errors.New("hi") + f := func(<-chan struct{}) error { + time.Sleep(10 * time.Millisecond) + return errVal + } + + wrapped := Run(f) + + // Once + { + err := wrapped.Wait() + if err != errVal { + t.Fatalf("bad: %#v", err) + } + } + + // Again + { + err := wrapped.Wait() + if err != errVal { + t.Fatalf("bad: %#v", err) + } + } +} + +func TestWrapped_cancel(t *testing.T) { + errVal := errors.New("hi") + f := func(ch <-chan struct{}) error { + <-ch + return errVal + } + + wrapped := Run(f) + + // Once + { + err := wrapped.Cancel() + if err != errVal { + t.Fatalf("bad: %#v", err) + } + } + + // Again + { + err := wrapped.Cancel() + if err != errVal { + t.Fatalf("bad: %#v", err) + } + } +}