diff --git a/signal/README.md b/signal/README.md new file mode 100644 index 00000000..e537ce1b --- /dev/null +++ b/signal/README.md @@ -0,0 +1,72 @@ +# Signal Server + +## Example + +```go +package main + +import ( + "os" + "syscall" + + "github.com/go-kratos/kratos/v2" + + "github.com/go-kratos-ecosystem/components/v2/signal" +) + +func main() { + app := kratos.New( + kratos.Server(newSignalServer()), + ) + + if err := app.Run(); err != nil { + panic(err) + } +} + +func newSignalServer() *signal.Server { + srv := signal.NewServer( + signal.WithRecoveryHandler(signal.DefaultRecoveryHandler), + ) + + srv.Register(&exampleHandler{}, &example2Handler{}) + + return srv +} + +type exampleHandler struct{} + +func (h *exampleHandler) Listen() []os.Signal { + return []os.Signal{syscall.SIGUSR1, syscall.SIGUSR2} +} + +func (h *exampleHandler) Handle(sig os.Signal) { + println("exampleHandler signal:", sig) +} + +type example2Handler struct{} + +func (h *example2Handler) Listen() []os.Signal { + return []os.Signal{syscall.SIGUSR1} +} + +func (h *example2Handler) Handle(os.Signal) { + panic("example2Handler panic") +} +``` + +Send signal: + +```bash +$ kill -SIGUSR2 42750 +$ kill -SIGUSR1 42750 +``` + +Output: + +```bash +INFO msg=[Signal] server starting +exampleHandler signal: (0x104ff0240,0x1051875b8) +exampleHandler signal: (0x104ff0240,0x1051875b0) +ERROR msg=[Signal] handler panic (user defined signal 1): example2Handler panic +``` \ No newline at end of file diff --git a/signal/handler.go b/signal/handler.go new file mode 100644 index 00000000..5820f402 --- /dev/null +++ b/signal/handler.go @@ -0,0 +1,10 @@ +package signal + +import ( + "os" +) + +type Handler interface { + Listen() []os.Signal + Handle(os.Signal) +} diff --git a/signal/server.go b/signal/server.go new file mode 100644 index 00000000..4191f9a5 --- /dev/null +++ b/signal/server.go @@ -0,0 +1,123 @@ +package signal + +import ( + "context" + "os" + "os/signal" + + "github.com/go-kratos/kratos/v2/log" +) + +var DefaultRecoveryHandler = func(err interface{}, sig os.Signal, handler Handler) { + log.Errorf("[Signal] handler panic (%s): %v", sig, err) +} + +type Server struct { + handlers []Handler + stoped chan struct{} + recoveryHandler func(interface{}, os.Signal, Handler) +} + +type Option func(*Server) + +func AddHandler(handler ...Handler) Option { + return func(s *Server) { + s.handlers = append(s.handlers, handler...) + } +} + +func WithRecoveryHandler(handler func(interface{}, os.Signal, Handler)) Option { + return func(s *Server) { + if handler != nil { + s.recoveryHandler = handler + } + } +} + +func NewServer(opts ...Option) *Server { + server := &Server{ + handlers: make([]Handler, 0), + stoped: make(chan struct{}), + } + + for _, opt := range opts { + opt(server) + } + + if server.recoveryHandler == nil { + server.recoveryHandler = DefaultRecoveryHandler + } + + return server +} + +func (s *Server) Start(ctx context.Context) error { + var ( + signals = make([]os.Signal, 0) + handlers = make(map[os.Signal][]Handler) + ) + + for _, h := range s.handlers { + for _, sig := range h.Listen() { + if _, ok := handlers[sig]; !ok { + handlers[sig] = make([]Handler, 0) + } + handlers[sig] = append(handlers[sig], h) + } + signals = append(signals, h.Listen()...) + } + + signals = s.uniqueSignals(signals) + + ch := make(chan os.Signal, len(signals)) + signal.Notify(ch, signals...) + + log.Infof("[Signal] server starting") + + for { + select { + case <-ctx.Done(): + return ctx.Err() + case <-s.stoped: + return nil + case sig := <-ch: + if hs, ok := handlers[sig]; ok { + for _, h := range hs { + s.handle(sig, h) + } + } + } + } +} + +func (s *Server) Register(handlers ...Handler) { + s.handlers = append(s.handlers, handlers...) +} + +func (s *Server) Stop(context.Context) error { + log.Infof("[Signal] server stopping") + close(s.stoped) + return nil +} + +func (s *Server) handle(sig os.Signal, handler Handler) { + defer func() { + if err := recover(); err != nil { + s.recoveryHandler(err, sig, handler) + } + }() + + handler.Handle(sig) +} + +func (s *Server) uniqueSignals(signals []os.Signal) []os.Signal { + m := make(map[os.Signal]struct{}) + for _, sig := range signals { + m[sig] = struct{}{} + } + signals = make([]os.Signal, 0) + for sig := range m { + signals = append(signals, sig) + } + return signals +} diff --git a/signal/server_test.go b/signal/server_test.go new file mode 100644 index 00000000..7a1c5a42 --- /dev/null +++ b/signal/server_test.go @@ -0,0 +1,77 @@ +package signal + +import ( + "bytes" + "context" + "fmt" + "os" + "sync" + "syscall" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +var ( + buffer bytes.Buffer + mu sync.Mutex +) + +func TestServer(t *testing.T) { + srv := newServer() + + go srv.Start(context.Background()) //nolint:errcheck + + time.Sleep(2 * time.Second) + + assert.NoError(t, syscall.Kill(os.Getpid(), syscall.SIGUSR1)) + assert.NoError(t, syscall.Kill(os.Getpid(), syscall.SIGUSR2)) + + time.Sleep(2 * time.Second) + + mu.Lock() + assert.Equal(t, `exampleHandler signal: user defined signal 1 +signal: user defined signal 1, handler: *signal.example2Handler, err: example2Handler panic +exampleHandler signal: user defined signal 2 +`, buffer.String()) + mu.Unlock() + + srv.Stop(context.Background()) //nolint:errcheck +} + +func newServer() *Server { + srv := NewServer( + WithRecoveryHandler(func(err interface{}, signal os.Signal, handler Handler) { + mu.Lock() + defer mu.Unlock() + buffer.WriteString(fmt.Sprintf("signal: %s, handler: %T, err: %v\n", signal, handler, err)) + }), + ) + + srv.Register(&exampleHandler{}, &example2Handler{}) + + return srv +} + +type exampleHandler struct{} + +func (h *exampleHandler) Listen() []os.Signal { + return []os.Signal{syscall.SIGUSR1, syscall.SIGUSR2} +} + +func (h *exampleHandler) Handle(sig os.Signal) { + mu.Lock() + defer mu.Unlock() + buffer.WriteString(fmt.Sprintf("exampleHandler signal: %s\n", sig)) +} + +type example2Handler struct{} + +func (h *example2Handler) Listen() []os.Signal { + return []os.Signal{syscall.SIGUSR1} +} + +func (h *example2Handler) Handle(os.Signal) { + panic("example2Handler panic") +}