Skip to content

Commit

Permalink
feat(signal): Added Signal Server (#76)
Browse files Browse the repository at this point in the history
* feat(signal): Added Signal Server

Signed-off-by: Flc゛ <[email protected]>

* feat(signal): Added Signal Server

Signed-off-by: Flc゛ <[email protected]>

* feat(signal): Added Signal Server

Signed-off-by: Flc゛ <[email protected]>

* feat(signal): Added Signal Server

Signed-off-by: Flc゛ <[email protected]>

---------

Signed-off-by: Flc゛ <[email protected]>
  • Loading branch information
flc1125 authored Jan 26, 2024
1 parent 99a279b commit 8011549
Show file tree
Hide file tree
Showing 4 changed files with 282 additions and 0 deletions.
72 changes: 72 additions & 0 deletions signal/README.md
Original file line number Diff line number Diff line change
@@ -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
```
10 changes: 10 additions & 0 deletions signal/handler.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,10 @@
package signal

import (
"os"
)

type Handler interface {
Listen() []os.Signal
Handle(os.Signal)
}
123 changes: 123 additions & 0 deletions signal/server.go
Original file line number Diff line number Diff line change
@@ -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
}
77 changes: 77 additions & 0 deletions signal/server_test.go
Original file line number Diff line number Diff line change
@@ -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")
}

0 comments on commit 8011549

Please sign in to comment.