feat: grace support custom handler

Author: mobus

examples/netepoll/main.go

@@ -0,0 +1,50 @@
+package main
+
+import (
+	"log"
+	"net/http"
+	_ "net/http/pprof"
+	"os"
+	"syscall"
+	"time"
+
+	"github.com/rcrowley/go-metrics"
+	"github.com/sunvim/utils/netpoll"
+)
+
+func main() {
+	setLimit()
+	go metrics.Log(metrics.DefaultRegistry, 5*time.Second, log.New(os.Stderr, "metrics: ", log.Lmicroseconds))
+
+	go func() {
+		if err := http.ListenAndServe(":6060", nil); err != nil {
+			log.Fatalf("pprof failed: %v", err)
+		}
+	}()
+
+	var handler = &netpoll.DataHandler{
+		NoShared:   true,
+		NoCopy:     true,
+		BufferSize: 1024,
+		HandlerFunc: func(req []byte) (res []byte) {
+			res = req
+			return
+		},
+	}
+	if err := netpoll.ListenAndServe("tcp", ":9999", handler); err != nil {
+		panic(err)
+	}
+}
+
+func setLimit() {
+	var rLimit syscall.Rlimit
+	if err := syscall.Getrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
+		panic(err)
+	}
+	rLimit.Cur = rLimit.Max
+	if err := syscall.Setrlimit(syscall.RLIMIT_NOFILE, &rLimit); err != nil {
+		panic(err)
+	}
+
+	log.Printf("set cur limit: %d", rLimit.Cur)
+}

go.mod

@@ -7,6 +7,7 @@ require (
 	github.com/hslam/scheduler v0.0.0-20211028175315-641598104976
 	github.com/hslam/sendfile v1.0.1
 	github.com/hslam/splice v1.0.3
+	github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475
 	github.com/stretchr/testify v1.7.0
 	gopkg.in/eapache/queue.v1 v1.1.0
 )

go.sum

@@ -12,6 +12,8 @@ github.com/hslam/splice v1.0.3 h1:CwSmzu6AAm8sb2wYgSGvTwixy3seqA6xJ9NXQ0ff3j4=
 github.com/hslam/splice v1.0.3/go.mod h1:7D1QlFptoG0ruXzcAwpzckKxUN4+ZpvrIhwfbcAQcx8=
 github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM=
 github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4=
+github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475 h1:N/ElC8H3+5XpJzTSTfLsJV/mx9Q9g7kxmchpfZyxgzM=
+github.com/rcrowley/go-metrics v0.0.0-20201227073835-cf1acfcdf475/go.mod h1:bCqnVzQkZxMG4s8nGwiZ5l3QUCyqpo9Y+/ZMZ9VjZe4=
 github.com/stretchr/objx v0.1.0/go.mod h1:HFkY916IF+rwdDfMAkV7OtwuqBVzrE8GR6GFx+wExME=
 github.com/stretchr/testify v1.7.0 h1:nwc3DEeHmmLAfoZucVR881uASk0Mfjw8xYJ99tb5CcY=
 github.com/stretchr/testify v1.7.0/go.mod h1:6Fq8oRcR53rry900zMqJjRRixrwX3KX962/h/Wwjteg=

grace/service.go

@@ -12,9 +12,11 @@ import (
 type serv struct {
 	stop   chan os.Signal
 	cancel func()
+	funcs  []func() error
 }
 
 type Service interface {
+	Register(fn func() error)
 	Wait()
 }
 
@@ -23,11 +25,20 @@ func (s *serv) Wait() {
 	select {
 	case <-s.stop:
 		s.cancel()
+		for _, fn := range s.funcs {
+			if err := fn(); err != nil {
+				log.Printf("err: %v \n", err)
+			}
+		}
 		time.Sleep(100 * time.Millisecond)
 		log.Println("all services exited totally.")
 	}
 }
 
+func (s *serv) Register(fn func() error) {
+	s.funcs = append(s.funcs, fn)
+}
+
 func New(ctx context.Context) (context.Context, Service) {
 	stopChan := make(chan os.Signal)
 	signal.Notify(stopChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)

grace/service_test.go

@@ -0,0 +1,26 @@
+package grace
+
+import (
+	"context"
+	"log"
+	"syscall"
+	"testing"
+	"time"
+)
+
+func TestWait(t *testing.T) {
+	_, service := New(context.Background())
+	service.Register(func() error {
+		log.Println("exit 1")
+		return nil
+	})
+	service.Register(func() error {
+		log.Println("exit 2")
+		return nil
+	})
+	go func() {
+		time.Sleep(10 * time.Second)
+		syscall.Kill(syscall.Getpid(), syscall.SIGINT)
+	}()
+	service.Wait()
+}