diff --git a/options.go b/options.go index 409f20f..cd9d682 100644 --- a/options.go +++ b/options.go @@ -76,12 +76,14 @@ type options struct { safeWrite bool maxUDPPacketSize int exactUDPBufferSizeEnabled bool + gracefulRestartTimeout time.Duration } func (o *options) setDefault() { o.tcpKeepAlive = defaultTCPKeepAlive o.maxUDPPacketSize = defaultUDPBufferSize o.exactUDPBufferSizeEnabled = defaultExactUDPBufferSizeEnabled + o.gracefulRestartTimeout = defaultGracefulRestartTimeout } // WithTCPKeepAlive sets the tcp keep alive interval. @@ -197,3 +199,11 @@ func WithExactUDPBufferSizeEnabled(exactUDPBufferSizeEnabled bool) Option { op.exactUDPBufferSizeEnabled = exactUDPBufferSizeEnabled }} } + +// WithGracefulRestartTimeout sets the timeout for graceful restart. +// The parent process waits this long after starting the child process before closing the listener. +func WithGracefulRestartTimeout(timeout time.Duration) Option { + return Option{func(op *options) { + op.gracefulRestartTimeout = timeout + }} +} diff --git a/restart_test.go b/restart_test.go new file mode 100644 index 0000000..b2b606f --- /dev/null +++ b/restart_test.go @@ -0,0 +1,255 @@ +// +// +// Tencent is pleased to support the open source community by making tRPC available. +// +// Copyright (C) 2023 Tencent. +// All rights reserved. +// +// If you have downloaded a copy of the tRPC source code from Tencent, +// please note that tRPC source code is licensed under the Apache 2.0 License, +// A copy of the Apache 2.0 License is included in this file. +// +// + +package tnet + +import ( + "context" + "io" + "net" + "os" + "os/exec" + "testing" + "time" + + "github.com/stretchr/testify/require" +) + +type serveOnlyService struct{} + +func (serveOnlyService) Serve(context.Context) error { return nil } + +var _ Service = serveOnlyService{} + +func TestCleanAndAppendEnv(t *testing.T) { + got := cleanAndAppendEnv([]string{ + "A=1", + gracefulRestartFDEnv + "=old", + "B=2", + }, gracefulRestartFDEnv, gracefulRestartFD) + require.Equal(t, []string{"A=1", "B=2", gracefulRestartFDEnv + "=" + gracefulRestartFD}, got) +} + +func TestWithGracefulRestartTimeout(t *testing.T) { + var opts options + opts.setDefault() + require.Equal(t, defaultGracefulRestartTimeout, opts.gracefulRestartTimeout) + + WithGracefulRestartTimeout(time.Millisecond).f(&opts) + require.Equal(t, time.Millisecond, opts.gracefulRestartTimeout) +} + +func TestServicesImplementRestartable(t *testing.T) { + ln, err := Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + tcpSvc, err := NewTCPService(ln, nil) + require.NoError(t, err) + require.Implements(t, (*Restartable)(nil), tcpSvc) + require.NoError(t, tcpSvc.(*tcpservice).close()) + + lns, err := ListenPackets("udp", "127.0.0.1:0", false) + require.NoError(t, err) + udpSvc, err := NewUDPService(lns, nil) + require.NoError(t, err) + require.Implements(t, (*Restartable)(nil), udpSvc) + require.NoError(t, udpSvc.(*udpservice).close()) +} + +func TestTCPServiceRestartStartError(t *testing.T) { + ln, err := Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + + svc, err := NewTCPService(ln, nil) + require.NoError(t, err) + ts := svc.(*tcpservice) + + origCommand := execCommand + t.Cleanup(func() { execCommand = origCommand }) + execCommand = func(string, ...string) *exec.Cmd { + return exec.Command("tnet-test-binary-does-not-exist") + } + + err = svc.(Restartable).Restart(context.Background()) + require.Error(t, err) + require.False(t, ts.restarting.Load()) +} + +func TestTCPServiceRestartKeepsActiveConnection(t *testing.T) { + ln, err := Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + + svc, err := NewTCPService(ln, echoTCP, WithGracefulRestartTimeout(10*time.Millisecond)) + require.NoError(t, err) + defer svc.(*tcpservice).close() + + lastCmd := mockRestartCommand(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + serveDone := serveTCP(t, svc, ctx) + client, err := net.Dial(ln.Addr().Network(), ln.Addr().String()) + require.NoError(t, err) + + assertTCPEcho(t, client, "before_restart") + + restartDone := make(chan error, 1) + go func() { + restartDone <- svc.(Restartable).Restart(context.Background()) + }() + + time.Sleep(50 * time.Millisecond) + select { + case err := <-restartDone: + t.Fatalf("restart returned before active connection drained: %v", err) + default: + } + + assertTCPEcho(t, client, "after_restart") + require.NoError(t, client.Close()) + + select { + case err := <-restartDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("restart did not finish after active connection closed") + } + requireRestartCommand(t, lastCmd()) + + select { + case err := <-serveDone: + require.NoError(t, err) + case <-time.After(time.Second): + t.Fatal("service did not return after restart drain") + } +} + +func TestTCPServiceWaitConnectionsContext(t *testing.T) { + ln, err := Listen("tcp", "127.0.0.1:0") + require.NoError(t, err) + defer ln.Close() + + svc, err := NewTCPService(ln, nil) + require.NoError(t, err) + ts := svc.(*tcpservice) + ts.conns[1001] = &tcpconn{nfd: netFD{fd: 1001}} + + ctx, cancel := context.WithCancel(context.Background()) + done := make(chan error, 1) + go func() { + done <- ts.waitConnections(ctx) + }() + time.Sleep(10 * time.Millisecond) + cancel() + + select { + case err := <-done: + require.ErrorIs(t, err, context.Canceled) + case <-time.After(time.Second): + t.Fatal("waitConnections did not return after context cancellation") + } + + ts.mu.Lock() + delete(ts.conns, 1001) + ts.mu.Unlock() +} + +func TestUDPServiceRestartClosesConns(t *testing.T) { + lns := make([]PacketConn, 0, 2) + for i := 0; i < 2; i++ { + rawConn, err := net.ListenPacket("udp", "127.0.0.1:0") + require.NoError(t, err) + conn, err := NewPacketConn(rawConn) + require.NoError(t, err) + lns = append(lns, conn) + } + t.Cleanup(func() { + for _, ln := range lns { + _ = ln.Close() + } + }) + + svc, err := newUDPService(lns, nil, WithGracefulRestartTimeout(0)) + require.NoError(t, err) + lastCmd := mockRestartCommand(t) + + err = svc.(Restartable).Restart(context.Background()) + require.NoError(t, err) + requireRestartCommand(t, lastCmd()) + for _, ln := range lns { + require.False(t, ln.IsActive()) + } +} + +func echoTCP(conn Conn) error { + buf := make([]byte, 1024) + n, err := conn.Read(buf) + if err != nil { + return err + } + _, err = conn.Write(buf[:n]) + return err +} + +func assertTCPEcho(t *testing.T, conn net.Conn, message string) { + t.Helper() + _, err := conn.Write([]byte(message)) + require.NoError(t, err) + require.NoError(t, conn.SetReadDeadline(time.Now().Add(time.Second))) + buf := make([]byte, len(message)) + _, err = io.ReadFull(conn, buf) + require.NoError(t, err) + require.Equal(t, message, string(buf)) +} + +func serveTCP(t *testing.T, svc Service, ctx context.Context) <-chan error { + t.Helper() + done := make(chan error, 1) + go func() { + done <- svc.Serve(ctx) + }() + time.Sleep(20 * time.Millisecond) + return done +} + +func mockRestartCommand(t *testing.T) func() *exec.Cmd { + t.Helper() + origCommand := execCommand + var last *exec.Cmd + t.Cleanup(func() { + execCommand = origCommand + if last != nil && last.Process != nil && last.ProcessState == nil { + _ = last.Wait() + } + }) + execCommand = func(string, ...string) *exec.Cmd { + last = exec.Command(os.Args[0], "-test.run=TestRestartHelperProcess", "--", "restart-helper") + return last + } + return func() *exec.Cmd { + return last + } +} + +func requireRestartCommand(t *testing.T, cmd *exec.Cmd) { + t.Helper() + require.NotNil(t, cmd) + require.Contains(t, cmd.Env, gracefulRestartFDEnv+"="+gracefulRestartFD) + require.Len(t, cmd.ExtraFiles, 1) +} + +func TestRestartHelperProcess(t *testing.T) { + if len(os.Args) > 1 && os.Args[len(os.Args)-1] == "restart-helper" { + os.Exit(0) + } +} diff --git a/tcpconn_outbound_test.go b/tcpconn_outbound_test.go index 273db54..c19de90 100644 --- a/tcpconn_outbound_test.go +++ b/tcpconn_outbound_test.go @@ -17,6 +17,7 @@ func (nopSockCloser) Close() error { func TestTCPConnWritevOutboundBufferLimitExceeded(t *testing.T) { tc := &tcpconn{ readTrigger: make(chan struct{}), + closedFinished: make(chan struct{}), outboundBufferLimit: 5, nfd: netFD{ sock: nopSockCloser{}, diff --git a/tcpservice.go b/tcpservice.go index a1f3cd6..2ad6458 100644 --- a/tcpservice.go +++ b/tcpservice.go @@ -18,6 +18,7 @@ import ( "errors" "fmt" "net" + "os" "sync" "time" @@ -56,6 +57,8 @@ func NewTCPService(listener net.Listener, handler TCPHandler, opt ...Option) (Se return newTCPService(ln, handler, opt...) } +var _ Restartable = (*tcpservice)(nil) + func newTCPService(ln *tcpListener, handler TCPHandler, opt ...Option) (Service, error) { opts := options{} opts.setDefault() @@ -70,6 +73,7 @@ func newTCPService(ln *tcpListener, handler TCPHandler, opt ...Option) (Service, conns: make(map[int]*tcpconn), hupCh: make(chan struct{}), } + s.connCond = sync.NewCond(&s.mu) return s, nil } @@ -80,14 +84,17 @@ const ( ) type tcpservice struct { - ln *tcpListener - reqHandle TCPHandler - hupCh chan struct{} - conns map[int]*tcpconn - opts options - closed atomic.Bool - tempDelay time.Duration - mu sync.Mutex + ln *tcpListener + reqHandle TCPHandler + hupCh chan struct{} + conns map[int]*tcpconn + opts options + closed atomic.Bool + restarting atomic.Bool + tempDelay time.Duration + mu sync.Mutex + hupOnce sync.Once + connCond *sync.Cond } // Serve starts the service. @@ -101,12 +108,15 @@ func (s *tcpservice) Serve(ctx context.Context) error { log.Infof("tnet tcp service started, current number of pollers: %d, use tnet.SetNumPollers to change it\n", poller.NumPollers()) - defer s.close() - select { case <-ctx.Done(): + _ = s.close() return ctx.Err() case <-s.hupCh: + if s.restarting.Load() { + return s.waitConnections(ctx) + } + _ = s.close() return errors.New("listener is closed") } } @@ -200,7 +210,9 @@ func tcpServiceOnHup(data interface{}) { if !ok || s == nil { panic(fmt.Sprintf("bug: data is not *tcpservice type (%v) or s is nil pointer (%v)", !ok, s == nil)) } - close(s.hupCh) + s.hupOnce.Do(func() { + close(s.hupCh) + }) } func (s *tcpservice) storeConn(conn *tcpconn) { @@ -213,11 +225,11 @@ func (s *tcpservice) storeConn(conn *tcpconn) { } func (s *tcpservice) deleteConn(conn *tcpconn) { - if s.closed.Load() { - return - } s.mu.Lock() - delete(s.conns, conn.nfd.FD()) + if _, ok := s.conns[conn.nfd.FD()]; ok { + delete(s.conns, conn.nfd.FD()) + s.connCond.Broadcast() + } s.mu.Unlock() } @@ -226,9 +238,70 @@ func (s *tcpservice) closeAll() { return } s.mu.Lock() + conns := make([]*tcpconn, 0, len(s.conns)) for k, conn := range s.conns { - conn.Close() + conns = append(conns, conn) delete(s.conns, k) } + s.connCond.Broadcast() s.mu.Unlock() + for _, conn := range conns { + conn.Close() + } +} + +func (s *tcpservice) waitConnections(ctx context.Context) error { + if ctx == nil { + ctx = context.Background() + } + + stop := make(chan struct{}) + go func() { + select { + case <-ctx.Done(): + s.mu.Lock() + s.connCond.Broadcast() + s.mu.Unlock() + case <-stop: + } + }() + defer close(stop) + + s.mu.Lock() + defer s.mu.Unlock() + for len(s.conns) > 0 { + if err := ctx.Err(); err != nil { + return err + } + s.connCond.Wait() + } + return ctx.Err() +} + +// Restart starts a new process, closes the listener, and waits for active TCP connections to drain. +func (s *tcpservice) Restart(ctx context.Context) error { + if s.closed.Load() { + return errors.New("service is closed") + } + if !s.restarting.CAS(false, true) { + return errors.New("service is already restarting") + } + + file := os.NewFile(uintptr(s.ln.FD()), gracefulListenerFileName) + cmd := execCommand(os.Args[0], os.Args[1:]...) + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + cmd.Env = cleanAndAppendEnv(os.Environ(), gracefulRestartFDEnv, gracefulRestartFD) + cmd.ExtraFiles = []*os.File{file} + if err := cmd.Start(); err != nil { + s.restarting.Store(false) + return err + } + + time.Sleep(s.opts.gracefulRestartTimeout) + if err := s.ln.Close(); err != nil { + return err + } + tcpServiceOnHup(s) + return s.waitConnections(ctx) } diff --git a/tnet.go b/tnet.go index ac49b6f..3ce3f7c 100644 --- a/tnet.go +++ b/tnet.go @@ -18,12 +18,35 @@ import ( "context" "fmt" "net" + "os/exec" + "strings" "time" // Stat reports user statistics. _ "trpc.group/trpc-go/tnet/internal/stat" ) +const ( + gracefulRestartFDEnv = "GRACEFUL_RESTART_FD" + gracefulRestartFD = "3" + gracefulListenerFileName = "listener" + gracefulUDPListenerFileName = "udp_listener" + defaultGracefulRestartTimeout = 2 * time.Second +) + +var execCommand = exec.Command + +func cleanAndAppendEnv(env []string, key, value string) []string { + prefix := key + "=" + out := make([]string, 0, len(env)+1) + for _, e := range env { + if !strings.HasPrefix(e, prefix) { + out = append(out, e) + } + } + return append(out, prefix+value) +} + // BaseConn is common for stream and packet oriented network connection. type BaseConn interface { // Conn extends net.Conn, just for interface compatibility. @@ -129,6 +152,12 @@ type Service interface { Serve(ctx context.Context) error } +// Restartable is optionally implemented by services that support graceful restart. +type Restartable interface { + // Restart starts a child process with inherited listener file descriptors and shuts down the current service. + Restart(ctx context.Context) error +} + // Listen announces on the local network address. // The network must be "tcp", "tcp4", "tcp6". func Listen(network, address string) (net.Listener, error) { diff --git a/udpservice.go b/udpservice.go index 57992f1..dcb9176 100644 --- a/udpservice.go +++ b/udpservice.go @@ -18,7 +18,9 @@ import ( "errors" "fmt" "net" + "os" "sync" + "time" goreuseport "github.com/kavu/go_reuseport" "trpc.group/trpc-go/tnet/internal/netutil" @@ -35,6 +37,8 @@ func NewUDPService(lns []PacketConn, handler UDPHandler, opt ...Option) (Service return newUDPService(lns, handler, opt...) } +var _ Restartable = (*udpservice)(nil) + func newUDPService(lns []PacketConn, handler UDPHandler, opt ...Option) (Service, error) { var opts options opts.setDefault() @@ -172,6 +176,25 @@ func (s *udpservice) close() error { return nil } +// Restart starts a new process with the first UDP listener fd and closes the old packet conns. +func (s *udpservice) Restart(ctx context.Context) error { + if len(s.conns) == 0 { + return errors.New("no UDP conn to restart") + } + file := os.NewFile(uintptr(s.conns[0].nfd.fd), gracefulUDPListenerFileName) + cmd := execCommand(os.Args[0], os.Args[1:]...) + cmd.Stdout = os.Stdout + cmd.Stderr = os.Stderr + cmd.Env = cleanAndAppendEnv(os.Environ(), gracefulRestartFDEnv, gracefulRestartFD) + cmd.ExtraFiles = []*os.File{file} + if err := cmd.Start(); err != nil { + return err + } + + time.Sleep(s.opts.gracefulRestartTimeout) + return s.close() +} + func validateListeners(lns []PacketConn) error { if len(lns) == 0 { return errors.New("listeners can't be nil")