github.com/angenalZZZ/gofunc@v0.0.0-20210507121333-48ff1be3917b/rpc/gracegrpc.go (about)

     1  // 1. In a terminal trigger a graceful server restart (using the pid from your output):
     2  //  > kill -USR2 pid
     3  // 2. Run with supervisor
     4  //  > command = /xxx/pid.py /xxx/log.pid /xxx/server
     5  //  > kill -USR2 {pid.py pid}
     6  
     7  package rpc
     8  
     9  import (
    10  	"fmt"
    11  	"github.com/angenalZZZ/gofunc/configfile"
    12  	"github.com/angenalZZZ/gofunc/log"
    13  	"net"
    14  	"os"
    15  	"os/signal"
    16  	"syscall"
    17  
    18  	"github.com/facebookgo/grace/gracenet"
    19  	"google.golang.org/grpc"
    20  )
    21  
    22  var (
    23  	SIGUSR2 = syscall.Signal(0x1f) // ReStart Process
    24  )
    25  
    26  // GraceGrpc is used to wrap a grpc server that can be gracefully terminated & restarted
    27  type GraceGrpc struct {
    28  	server   *grpc.Server
    29  	grace    *gracenet.Net
    30  	listener net.Listener
    31  	errors   chan error
    32  	pidPath  string
    33  	*log.Logger
    34  }
    35  
    36  // NewGraceGrpc is used to construct a new GraceGrpc
    37  func NewGraceGrpc(s *grpc.Server, net, addr, pidPath string, logCfgFile string) (*GraceGrpc, error) {
    38  	if logCfgFile == "" {
    39  		logCfgFile = "log.yaml"
    40  	}
    41  	logCfg := new(log.AConfig)
    42  	if err := configfile.YamlTo(logCfgFile, logCfg); err != nil {
    43  		_ = fmt.Errorf("%s\n", err.Error())
    44  	}
    45  
    46  	gr := &GraceGrpc{
    47  		server:  s,
    48  		grace:   &gracenet.Net{},
    49  		errors:  make(chan error),
    50  		pidPath: pidPath,
    51  		Logger:  log.Init(logCfg.Log),
    52  	}
    53  	listener, err := gr.grace.Listen(net, addr)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  	gr.listener = listener
    58  	return gr, nil
    59  }
    60  
    61  func (gr *GraceGrpc) startServe() {
    62  	if err := gr.server.Serve(gr.listener); err != nil {
    63  		gr.errors <- err
    64  	}
    65  }
    66  
    67  func (gr *GraceGrpc) handleSignal() <-chan struct{} {
    68  	terminate := make(chan struct{})
    69  	go func() {
    70  		ch := make(chan os.Signal, 10)
    71  		signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM, SIGUSR2)
    72  		for {
    73  			sig := <-ch
    74  			switch sig {
    75  			case syscall.SIGINT, syscall.SIGTERM:
    76  				signal.Stop(ch)
    77  				gr.server.GracefulStop()
    78  				close(terminate)
    79  				return
    80  			case SIGUSR2:
    81  				if _, err := gr.grace.StartProcess(); err != nil {
    82  					gr.errors <- err
    83  				}
    84  			}
    85  		}
    86  	}()
    87  	return terminate
    88  }
    89  
    90  // storePid is used to write out PID to pidPath
    91  func (gr *GraceGrpc) storePid(pid int) error {
    92  	pidPath := gr.pidPath
    93  	if pidPath == "" {
    94  		return fmt.Errorf("no pid file path: %s", pidPath)
    95  	}
    96  
    97  	pidFile, err := os.OpenFile(pidPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
    98  	if err != nil {
    99  		return fmt.Errorf("could not open pid file: %v", err)
   100  	}
   101  	defer pidFile.Close()
   102  
   103  	_, err = pidFile.WriteString(fmt.Sprintf("%d", pid))
   104  	if err != nil {
   105  		return fmt.Errorf("could not write to pid file: %s", err)
   106  	}
   107  	return nil
   108  }
   109  
   110  // Serve is used to start grpc server.
   111  // Serve will gracefully terminated or restarted when handling signals.
   112  func (gr *GraceGrpc) Serve() error {
   113  	if gr.listener == nil {
   114  		return fmt.Errorf("gracegrpc must construct by new\n")
   115  	}
   116  
   117  	pid := os.Getpid()
   118  	addrString := gr.listener.Addr().String()
   119  
   120  	gr.Info().Msgf("Serving %s with pid %d\n", addrString, pid)
   121  
   122  	if err := gr.storePid(pid); err != nil {
   123  		return err
   124  	}
   125  
   126  	go gr.startServe()
   127  
   128  	terminate := gr.handleSignal()
   129  
   130  	select {
   131  	case err := <-gr.errors:
   132  		return err
   133  	case <-terminate:
   134  		gr.Error().Msgf("Exiting pid %d", os.Getpid())
   135  		return nil
   136  	}
   137  }