github.com/sunvim/utils@v0.1.0/grace/service.go (about)

     1  package grace
     2  
     3  import (
     4  	"context"
     5  	"log"
     6  	"os"
     7  	"os/signal"
     8  	"syscall"
     9  	"time"
    10  )
    11  
    12  type serv struct {
    13  	ctx    context.Context
    14  	stop   chan os.Signal
    15  	cancel func()
    16  	funcs  []func() error
    17  	srvs   map[string]func(context.Context) error
    18  }
    19  
    20  type Service interface {
    21  	Register(fn func() error)
    22  	RegisterService(name string, fn func(context.Context) error)
    23  	Wait()
    24  }
    25  
    26  func (s *serv) Wait() {
    27  	defer signal.Stop(s.stop)
    28  	for name, fn := range s.srvs {
    29  		log.Printf("boot service %s ...", name)
    30  		safeGo(s.ctx, fn)
    31  	}
    32  	select {
    33  	case <-s.stop:
    34  		s.cancel()
    35  		for _, fn := range s.funcs {
    36  			if err := fn(); err != nil {
    37  				log.Printf("err: %v \n", err)
    38  			}
    39  		}
    40  		time.Sleep(100 * time.Millisecond)
    41  		log.Println("all services exited totally.")
    42  		os.Exit(0)
    43  	}
    44  }
    45  
    46  func (s *serv) Register(fn func() error) {
    47  	s.funcs = append(s.funcs, fn)
    48  }
    49  
    50  func (s *serv) RegisterService(name string, fn func(context.Context) error) {
    51  	s.srvs[name] = fn
    52  }
    53  
    54  func New(ctx context.Context) (context.Context, Service) {
    55  	stopChan := make(chan os.Signal)
    56  	signal.Notify(stopChan, syscall.SIGHUP, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
    57  	ctxLocal, cancel := context.WithCancel(ctx)
    58  	return ctxLocal, &serv{
    59  		ctx:    ctxLocal,
    60  		stop:   stopChan,
    61  		cancel: cancel,
    62  		srvs:   make(map[string]func(context.Context) error),
    63  	}
    64  }
    65  
    66  func safeGo(ctx context.Context, fn func(ctx context.Context) error) {
    67  	go func(ctx context.Context) {
    68  		defer func() {
    69  			if err := recover(); err != nil {
    70  				println(err)
    71  			}
    72  		}()
    73  		fn(ctx)
    74  	}(ctx)
    75  }