Skip to content

Commit bdcea55

Browse files
committed
helper/signalwrapper
1 parent 0c8b243 commit bdcea55

2 files changed

Lines changed: 245 additions & 0 deletions

File tree

‎helper/signalwrapper/wrapper.go‎

Lines changed: 177 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,177 @@
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+
}
Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,68 @@
1+
package signalwrapper
2+
3+
import (
4+
"errors"
5+
"testing"
6+
"time"
7+
)
8+
9+
func TestWrapped_goodCh(t *testing.T) {
10+
errVal := errors.New("hi")
11+
f := func(<-chan struct{}) error { return errVal }
12+
err := <-Run(f).ErrCh
13+
if err != errVal {
14+
t.Fatalf("bad: %#v", err)
15+
}
16+
}
17+
18+
func TestWrapped_goodWait(t *testing.T) {
19+
errVal := errors.New("hi")
20+
f := func(<-chan struct{}) error {
21+
time.Sleep(10 * time.Millisecond)
22+
return errVal
23+
}
24+
25+
wrapped := Run(f)
26+
27+
// Once
28+
{
29+
err := wrapped.Wait()
30+
if err != errVal {
31+
t.Fatalf("bad: %#v", err)
32+
}
33+
}
34+
35+
// Again
36+
{
37+
err := wrapped.Wait()
38+
if err != errVal {
39+
t.Fatalf("bad: %#v", err)
40+
}
41+
}
42+
}
43+
44+
func TestWrapped_cancel(t *testing.T) {
45+
errVal := errors.New("hi")
46+
f := func(ch <-chan struct{}) error {
47+
<-ch
48+
return errVal
49+
}
50+
51+
wrapped := Run(f)
52+
53+
// Once
54+
{
55+
err := wrapped.Cancel()
56+
if err != errVal {
57+
t.Fatalf("bad: %#v", err)
58+
}
59+
}
60+
61+
// Again
62+
{
63+
err := wrapped.Cancel()
64+
if err != errVal {
65+
t.Fatalf("bad: %#v", err)
66+
}
67+
}
68+
}

0 commit comments

Comments
 (0)