ADD: osutil.WaitForStop()

util/osutil/osutil.go

@@ -4,13 +4,24 @@ package osutil
 import (
 	"os"
 	"os/signal"
+	"syscall"
 )
 
-// WaitForInterrupt blocks until current process receives SIGINT.
-func WaitForInterrupt() {
+// WaitFor blocks until current process receives any of signals.
+func WaitFor(signals ...os.Signal) {
 	c := make(chan os.Signal, 1)
-	signal.Notify(c, os.Interrupt)
+	signal.Notify(c, signals...)
 	defer signal.Stop(c)
 
 	<-c
 }
+
+// WaitForInterrupt blocks until current process receives SIGINT.
+func WaitForInterrupt() {
+	WaitFor(syscall.SIGINT)
+}
+
+// WaitForStop blocks until current process receives SIGINT or SIGTERM.
+func WaitForStop() {
+	WaitFor(syscall.SIGINT, syscall.SIGTERM)
+}

util/osutil/osutil_test.go

@@ -3,21 +3,24 @@ package osutil_test
 import (
 	"os"
 	"runtime"
+	"syscall"
 	"testing"
 	"time"
 
 	"github.com/whoisnian/glb/util/osutil"
 )
 
+const (
+	waitingTime = 50 * time.Millisecond
+	timeoutTime = 60 * time.Millisecond
+)
+
 func TestWaitForInterrupt(t *testing.T) {
 	if runtime.GOOS == "windows" {
-		t.Skipf("sending os.Interrupt is not implemented for windows")
+		t.Skipf("sending syscall.SIGINT is not implemented for windows")
 	}
 
 	done := make(chan struct{})
-	waitingTime := 50 * time.Millisecond
-	timeoutTime := 10*time.Millisecond + waitingTime
-
 	go func() {
 		osutil.WaitForInterrupt()
 		close(done)
@@ -27,11 +30,35 @@ func TestWaitForInterrupt(t *testing.T) {
 	if err != nil {
 		t.Fatalf("os.FindProcess() error: %v", err)
 	}
-	time.AfterFunc(waitingTime, func() { p.Signal(os.Interrupt) })
 
+	time.AfterFunc(waitingTime, func() { p.Signal(syscall.SIGINT) })
 	select {
 	case <-done:
 	case <-time.After(timeoutTime):
 		t.Fatal("WaitForInterrupt timeout")
 	}
 }
+
+func TestWaitForStop(t *testing.T) {
+	if runtime.GOOS == "windows" {
+		t.Skipf("sending syscall.SIGTERM is not implemented for windows")
+	}
+
+	done := make(chan struct{})
+	go func() {
+		osutil.WaitForStop()
+		close(done)
+	}()
+
+	p, err := os.FindProcess(os.Getpid())
+	if err != nil {
+		t.Fatalf("os.FindProcess() error: %v", err)
+	}
+
+	time.AfterFunc(waitingTime, func() { p.Signal(syscall.SIGTERM) })
+	select {
+	case <-done:
+	case <-time.After(timeoutTime):
+		t.Fatal("WaitForStop timeout")
+	}
+}