Skip to content

Commit 0ddf69b

Browse files
committed
Add termination package
1 parent ca94683 commit 0ddf69b

File tree

2 files changed

+78
-0
lines changed

2 files changed

+78
-0
lines changed
Lines changed: 42 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,42 @@
1+
/*
2+
Package termination deals with Ctrl+C termination.
3+
*/
4+
package termination
5+
6+
import (
7+
"context"
8+
"fmt"
9+
"io"
10+
"os"
11+
"os/signal"
12+
"syscall"
13+
)
14+
15+
var cleanFns []func()
16+
17+
// for testing
18+
var osExit = os.Exit
19+
20+
// Listen listens signals.
21+
func Listen(ctx context.Context, w io.Writer) (context.Context, func()) {
22+
ctx, cancel := context.WithCancel(ctx)
23+
24+
ch := make(chan os.Signal, 2)
25+
signal.Notify(ch, os.Interrupt, syscall.SIGTERM)
26+
go func() {
27+
<-ch
28+
fmt.Fprintln(w, "\rCtrl+C pressed in Terminal")
29+
cancel()
30+
for _, f := range cleanFns {
31+
f()
32+
}
33+
osExit(0)
34+
}()
35+
36+
return ctx, cancel
37+
}
38+
39+
// CleanFunc registers clean function.
40+
func CleanFunc(f func()) {
41+
cleanFns = append(cleanFns, f)
42+
}
Lines changed: 36 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,36 @@
1+
package termination
2+
3+
import (
4+
"context"
5+
"io/ioutil"
6+
"os"
7+
"testing"
8+
"time"
9+
)
10+
11+
func TestTermination_Listen(t *testing.T) {
12+
CleanFunc(func() {})
13+
14+
doneCh := make(chan struct{})
15+
osExit = func(code int) { doneCh <- struct{}{} }
16+
17+
_, clean := Listen(context.Background(), ioutil.Discard)
18+
defer clean()
19+
20+
proc, err := os.FindProcess(os.Getpid())
21+
if err != nil {
22+
t.Fatalf("err %s", err)
23+
}
24+
25+
err = proc.Signal(os.Interrupt)
26+
if err != nil {
27+
t.Fatalf("err %s", err)
28+
}
29+
30+
select {
31+
case <-doneCh:
32+
return
33+
case <-time.After(100 * time.Millisecond):
34+
t.Fatal("timeout")
35+
}
36+
}

0 commit comments

Comments
 (0)