|
| 1 | +// Package signalwrapper is used to run functions that are sensitive to |
| 2 | +// signals that may be received from outside the process. It can also be |
| 3 | +// used as just an async function runner that is cancellable and we may |
| 4 | +// abstract this further into another package in the future. |
| 5 | +package signalwrapper |
| 6 | + |
| 7 | +import ( |
| 8 | + "os" |
| 9 | + "os/signal" |
| 10 | + "sync" |
| 11 | +) |
| 12 | + |
| 13 | +// CancellableFunc is a function that cancels if it receives a message |
| 14 | +// on the given channel. It must return an error if any occurred. It should |
| 15 | +// return no error if it was cancelled successfully since it is assumed |
| 16 | +// that this function will probably be called again at some future point |
| 17 | +// since it was interrupted. |
| 18 | +type CancellableFunc func(<-chan struct{}) error |
| 19 | + |
| 20 | +// Run wraps and runs the given cancellable function and returns the Wrapped |
| 21 | +// struct that can be used to listen for events, cancel on other events |
| 22 | +// (such as timeouts), etc. |
| 23 | +func Run(f CancellableFunc) *Wrapped { |
| 24 | + // Determine the signals we're listening to. Prematurely making |
| 25 | + // this a slice since I predict a future where we'll add others and |
| 26 | + // the complexity in doing so is low. |
| 27 | + signals := []os.Signal{os.Interrupt} |
| 28 | + |
| 29 | + // Register a listener for the signals |
| 30 | + sigCh := make(chan os.Signal, 1) |
| 31 | + signal.Notify(sigCh, signals...) |
| 32 | + |
| 33 | + // Create the channel we'll use to "cancel" |
| 34 | + cancelCh := make(chan struct{}) |
| 35 | + |
| 36 | + // This channel keeps track of whether the function we're running |
| 37 | + // completed successfully and the errors it may have had. It is |
| 38 | + // VERY IMPORTANT that the errCh is buffered to at least 1 so that |
| 39 | + // it doesn't block when finishing. |
| 40 | + doneCh := make(chan struct{}) |
| 41 | + errCh := make(chan error, 1) |
| 42 | + |
| 43 | + // Build our wrapped result |
| 44 | + wrapped := &Wrapped{ |
| 45 | + ErrCh: errCh, |
| 46 | + errCh: errCh, |
| 47 | + cancelCh: cancelCh, |
| 48 | + } |
| 49 | + |
| 50 | + // Start the function |
| 51 | + go func() { |
| 52 | + err := f(cancelCh) |
| 53 | + |
| 54 | + // Close the done channel _before_ sending the error in case |
| 55 | + // the error channel read blocks (it shouldn't) to avoid interrupts |
| 56 | + // doing anything. |
| 57 | + close(doneCh) |
| 58 | + |
| 59 | + // Mark completion |
| 60 | + wrapped.done(err) |
| 61 | + }() |
| 62 | + |
| 63 | + // Goroutine to track interrupts and make sure we do at-most-once |
| 64 | + // delivery of an interrupt since we're using a channel. |
| 65 | + go func() { |
| 66 | + // Clean up after this since this is the only function that |
| 67 | + // reads signals. |
| 68 | + defer signal.Stop(sigCh) |
| 69 | + |
| 70 | + select { |
| 71 | + case <-doneCh: |
| 72 | + // Everything happened naturally |
| 73 | + case <-sigCh: |
| 74 | + // Stop the function. Goroutine since we don't care about |
| 75 | + // the result and we'd like to end this goroutine as soon |
| 76 | + // as possible to avoid any more signals coming in. |
| 77 | + go wrapped.Cancel() |
| 78 | + } |
| 79 | + }() |
| 80 | + |
| 81 | + return wrapped |
| 82 | +} |
| 83 | + |
| 84 | +// Wrapped is the return value of wrapping a function. This has channels |
| 85 | +// that can be used to wait for a result as well as functions to help with |
| 86 | +// different behaviors. |
| 87 | +type Wrapped struct { |
| 88 | + // Set and consumed by user |
| 89 | + |
| 90 | + // ErrCh is the channel to listen for real-time events on the wrapped |
| 91 | + // function. A nil error sent means the execution completed without error. |
| 92 | + // This is an exactly once delivery channel. |
| 93 | + ErrCh <-chan error |
| 94 | + |
| 95 | + // Set by creator |
| 96 | + errCh chan<- error |
| 97 | + cancelCh chan<- struct{} |
| 98 | + |
| 99 | + // Set automatically |
| 100 | + once sync.Once |
| 101 | + cancelCond *sync.Cond |
| 102 | + cancelLock *sync.Mutex |
| 103 | + resultErr error |
| 104 | + resultSet bool |
| 105 | +} |
| 106 | + |
| 107 | +// Cancel stops the running function and blocks until it returns. The |
| 108 | +// resulting value is returned. |
| 109 | +// |
| 110 | +// It is safe to call this multiple times. This will return the resulting |
| 111 | +// error value each time. |
| 112 | +func (w *Wrapped) Cancel() error { |
| 113 | + w.once.Do(w.init) |
| 114 | + w.cancelLock.Lock() |
| 115 | + |
| 116 | + // If we have a result set, return that |
| 117 | + if w.resultSet { |
| 118 | + w.cancelLock.Unlock() |
| 119 | + return w.resultErr |
| 120 | + } |
| 121 | + |
| 122 | + // If we have a cancel channel, close it to signal and set it to |
| 123 | + // nil so we never do that again. |
| 124 | + if w.cancelCh != nil { |
| 125 | + close(w.cancelCh) |
| 126 | + w.cancelCh = nil |
| 127 | + } |
| 128 | + |
| 129 | + // Wait for the result to be set |
| 130 | + defer w.cancelLock.Unlock() |
| 131 | + w.cancelCond.Wait() |
| 132 | + return w.resultErr |
| 133 | +} |
| 134 | + |
| 135 | +// Wait waits for the completion of the wrapped function and returns the result. |
| 136 | +// |
| 137 | +// This can be called multiple times safely. |
| 138 | +func (w *Wrapped) Wait() error { |
| 139 | + w.once.Do(w.init) |
| 140 | + w.cancelLock.Lock() |
| 141 | + defer w.cancelLock.Unlock() |
| 142 | + |
| 143 | + // If we don't have a result yet, wait for that |
| 144 | + if !w.resultSet { |
| 145 | + w.cancelCond.Wait() |
| 146 | + } |
| 147 | + |
| 148 | + // Return the result |
| 149 | + return w.resultErr |
| 150 | +} |
| 151 | + |
| 152 | +// done marks this wrapped function as done with the resulting value. |
| 153 | +// This must only be called once. |
| 154 | +func (w *Wrapped) done(err error) { |
| 155 | + w.once.Do(w.init) |
| 156 | + w.cancelLock.Lock() |
| 157 | + |
| 158 | + // Set the result |
| 159 | + w.resultErr = err |
| 160 | + w.resultSet = true |
| 161 | + |
| 162 | + // Notify any waiters |
| 163 | + w.cancelCond.Broadcast() |
| 164 | + |
| 165 | + // Unlock since the next call can be blocking |
| 166 | + w.cancelLock.Unlock() |
| 167 | + |
| 168 | + // Notify any channel listeners |
| 169 | + w.errCh <- err |
| 170 | +} |
| 171 | + |
| 172 | +func (w *Wrapped) init() { |
| 173 | + // Create the condition variable |
| 174 | + var m sync.Mutex |
| 175 | + w.cancelCond = sync.NewCond(&m) |
| 176 | + w.cancelLock = &m |
| 177 | +} |
0 commit comments