package tcp /** * A tcp server */ import ( "context" "fmt" "github.com/hdt3213/godis/interface/tcp" "github.com/hdt3213/godis/lib/logger" "github.com/hdt3213/godis/lib/sync/atomic" "net" "os" "os/signal" "sync" "syscall" "time" ) // Config stores tcp server properties type Config struct { Address string `yaml:"address"` MaxConnect uint32 `yaml:"max-connect"` Timeout time.Duration `yaml:"timeout"` } // ListenAndServeWithSignal binds port and handle requests, blocking until receive stop signal func ListenAndServeWithSignal(cfg *Config, handler tcp.Handler) error { closeChan := make(chan struct{}) sigCh := make(chan os.Signal) signal.Notify(sigCh, syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGINT) go func() { sig := <-sigCh switch sig { case syscall.SIGHUP, syscall.SIGQUIT, syscall.SIGTERM, syscall.SIGINT: closeChan <- struct{}{} } }() listener, err := net.Listen("tcp", cfg.Address) if err != nil { return err } //cfg.Address = listener.Addr().String() logger.Info(fmt.Sprintf("bind: %s, start listening...", cfg.Address)) ListenAndServe(listener, handler, closeChan) return nil } // ListenAndServe binds port and handle requests, blocking until close func ListenAndServe(listener net.Listener, handler tcp.Handler, closeChan <-chan struct{}) { // listen signal var closing atomic.Boolean go func() { <-closeChan logger.Info("shutting down...") closing.Set(true) _ = listener.Close() // listener.Accept() will return err immediately _ = handler.Close() // close connections }() // listen port defer func() { // close during unexpected error _ = listener.Close() _ = handler.Close() }() ctx := context.Background() var waitDone sync.WaitGroup for { conn, err := listener.Accept() if err != nil { if closing.Get() { logger.Info("waiting disconnect...") waitDone.Wait() return // handler will be closed by defer } logger.Error(fmt.Sprintf("accept err: %v", err)) continue } // handle logger.Info("accept link") waitDone.Add(1) go func() { defer func() { waitDone.Done() }() handler.Handle(ctx, conn) }() } }