feat: listen and shutdown (#4)

* feat: ws handler shutdown

* fix

Author: zijiren233

client/wsd.go

@@ -273,11 +273,3 @@ func (wc *Dialer) DialContext(ctx context.Context, options ...ConnectOption) (ne
 func (wc *Dialer) Dial(options ...ConnectOption) (net.Conn, error) {
 	return wc.DialContext(context.Background(), options...)
 }
-
-func (wc *Dialer) DialTCP(options ...ConnectOption) (net.Conn, error) {
-	return wc.Dial(options...)
-}
-
-func (wc *Dialer) DialContextTCP(ctx context.Context, options ...ConnectOption) (net.Conn, error) {
-	return wc.DialContext(ctx, options...)
-}

server/healthy.go

@@ -7,7 +7,6 @@ import (
 	"log"
 	"net/http"
 	"os"
-	"sync/atomic"
 	"time"
 )
 
@@ -17,9 +16,11 @@ var (
 	jwtToken      = os.Getenv("JWT_TOKEN")
 )
 
-var ActiveNum int32 = 0
+type ActiveNum interface {
+	ActiveNum() int64
+}
 
-func HealthyCheck(ctx context.Context) {
+func HealthyCheck(ctx context.Context, handler ActiveNum) {
 	shutdownDuration, _ := time.ParseDuration(interval)
 	ticker := time.NewTicker(1 * time.Minute)
 	defer ticker.Stop()
@@ -28,7 +29,7 @@ func HealthyCheck(ctx context.Context) {
 	for {
 		select {
 		case <-ticker.C:
-			if atomic.LoadInt32(&ActiveNum) == 0 {
+			if handler.ActiveNum() == 0 {
 				zeroDuration += 1 * time.Minute
 				if zeroDuration >= shutdownDuration {
 					sendShutdownRequest()

server/main.go

@@ -3,10 +3,12 @@ package main
 import (
 	"context"
 	"log"
+	"net"
 	"os"
 	"os/signal"
 	"sync"
 	"syscall"
+	"time"
 )
 
 var (
@@ -16,29 +18,39 @@ var (
 )
 
 func main() {
-	ctx, cancel := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
-	defer cancel()
+	if listen == "" || target == "" {
+		log.Fatalf("LISTEN or TARGET is not set")
+	}
+
+	listener, err := net.Listen("tcp", listen)
+	if err != nil {
+		log.Fatalf("Failed to listen: %v", err)
+	}
+
+	ctx, stop := signal.NotifyContext(context.Background(), syscall.SIGINT, syscall.SIGTERM)
+	defer stop()
 
 	var wg sync.WaitGroup
 
+	handler := NewHandler(target)
+
 	if flag == "true" {
 		wg.Add(1)
 		go func() {
 			defer wg.Done()
-			HealthyCheck(ctx)
+			HealthyCheck(ctx, handler)
 		}()
 	}
 
-	if listen == "" || target == "" {
-		panic("LISTEN or TARGET is not set")
-	}
+	server := NewServer(
+		"/",
+		handler,
+		WithListener(listener),
+	)
 
 	go func() {
-		err := NewServer(
-			listen,
-			"/",
-			NewHandler(target),
-		).Serve()
+		log.Println("Started ws server")
+		err := server.Serve()
 		if err != nil {
 			log.Fatalf("Failed to ws serve: %v", err)
 		}
@@ -47,4 +59,11 @@ func main() {
 	<-ctx.Done()
 	log.Println("Shutting down gracefully")
 	wg.Wait()
+
+	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+	defer cancel()
+	err = server.Shutdown(ctx)
+	if err != nil {
+		log.Fatalf("Failed to shutdown gracefully: %v", err)
+	}
 }

server/ws.go

@@ -9,6 +9,7 @@ import (
 )
 
 type Server struct {
+	listener          net.Listener
 	listenErr         error
 	shutdowned        chan struct{}
 	onListened        chan struct{}
@@ -21,9 +22,20 @@ type Server struct {
 
 type ServerOption func(*Server)
 
-func NewServer(listenAddr, path string, wsHandler *Handler, opts ...ServerOption) *Server {
+func WithListener(listener net.Listener) ServerOption {
+	return func(ps *Server) {
+		ps.listener = listener
+	}
+}
+
+func WithListenAddr(listenAddr string) ServerOption {
+	return func(ps *Server) {
+		ps.listenAddr = listenAddr
+	}
+}
+
+func NewServer(path string, wsHandler *Handler, opts ...ServerOption) *Server {
 	ps := &Server{
-		listenAddr: listenAddr,
 		wsHandler:  wsHandler,
 		path:       path,
 		onListened: make(chan struct{}),
@@ -43,25 +55,13 @@ func (ps *Server) closeOnListened() {
 	})
 }
 
-func (ps *Server) OnListened() <-chan struct{} {
-	return ps.onListened
-}
-
-func (ps *Server) ListenErr() error {
+func (ps *Server) WaitListen() error {
+	<-ps.onListened
 	return ps.listenErr
 }
 
-func (ps *Server) Shutdowned() <-chan struct{} {
-	return ps.shutdowned
-}
-
-func (ps *Server) ShutdownedBool() bool {
-	select {
-	case <-ps.shutdowned:
-		return true
-	default:
-		return false
-	}
+func (ps *Server) WaitShutdown() {
+	<-ps.shutdowned
 }
 
 func (ps *Server) Serve() error {
@@ -73,12 +73,19 @@ func (ps *Server) Serve() error {
 	return ps.listenAndServe(server)
 }
 
-func (ps *Server) listenAndServe(server *http.Server) error {
+func (ps *Server) getListener() (net.Listener, error) {
+	if ps.listener != nil {
+		return ps.listener, nil
+	}
 	addr := ps.listenAddr
 	if addr == "" {
 		addr = ":http"
 	}
-	ln, err := net.Listen("tcp", addr)
+	return net.Listen("tcp", addr)
+}
+
+func (ps *Server) listenAndServe(server *http.Server) error {
+	ln, err := ps.getListener()
 	if err != nil {
 		ps.listenErr = err
 		return err
@@ -100,17 +107,20 @@ func (ps *Server) Server() *http.Server {
 			ReadHeaderTimeout: time.Second * 5,
 			MaxHeaderBytes:    16 * 1024,
 		}
+		ps.server.RegisterOnShutdown(func() {
+			ps.wsHandler.Close()
+		})
 	}
 	return ps.server
 }
 
 func (ps *Server) Close() error {
-	timeoutCtx, cancel := context.WithTimeout(context.Background(), 3*time.Second)
-	defer cancel()
-	return ps.Shutdown(timeoutCtx)
+	defer ps.closeOnListened()
+	return ps.server.Close()
 }
 
 func (ps *Server) Shutdown(ctx context.Context) error {
-	ps.closeOnListened()
+	defer ps.closeOnListened()
+	defer ps.wsHandler.Wait()
 	return ps.server.Shutdown(ctx)
 }

server/wsh.go

@@ -42,8 +42,12 @@ type GetTargetFunc func(req *http.Request) (string, []string, error)
 type Handler struct {
 	bufferPool        *sync.Pool
 	wsServer          *websocket.Server
+	closeChan         chan struct{}
 	defaultTargetAddr string
+	connectionsWg     sync.WaitGroup
+	activeNum         int64
 	bufferSize        int
+	closeOnce         sync.Once
 }
 
 type HandlerOption func(*Handler)
@@ -65,6 +69,7 @@ func checkOrigin(config *websocket.Config, req *http.Request) (err error) {
 func NewHandler(targetAddr string, opts ...HandlerOption) *Handler {
 	h := &Handler{
 		defaultTargetAddr: targetAddr,
+		closeChan:         make(chan struct{}),
 	}
 
 	for _, opt := range opts {
@@ -95,10 +100,26 @@ func (h *Handler) putBuffer(buffer *[]byte) {
 	}
 }
 
+func (h *Handler) ActiveNum() int64 {
+	return atomic.LoadInt64(&h.activeNum)
+}
+
+func (h *Handler) addActiveNum() {
+	atomic.AddInt64(&h.activeNum, 1)
+}
+
+func (h *Handler) subActiveNum() {
+	atomic.AddInt64(&h.activeNum, -1)
+}
+
 func (h *Handler) ServeHTTP(w http.ResponseWriter, req *http.Request) {
-	atomic.AddInt32(&ActiveNum, 1)
+	h.connectionsWg.Add(1)
+	defer h.connectionsWg.Done()
+
+	h.addActiveNum()
+	defer h.subActiveNum()
+
 	h.wsServer.ServeHTTP(w, req)
-	atomic.AddInt32(&ActiveNum, -1)
 }
 
 var pingCodec = websocket.Codec{
@@ -127,6 +148,9 @@ func (h *Handler) handleWebSocket(ws *websocket.Conn) {
 				}
 				_ = ws.Close()
 				return
+			case <-h.closeChan:
+				_ = ws.Close()
+				return
 			case <-exit:
 				return
 			}
@@ -197,3 +221,13 @@ func CopyBufferWithWriteTimeout(dst deadlineWriter, src io.Reader, buf []byte, t
 	}
 	return written, err
 }
+
+func (h *Handler) Close() {
+	h.closeOnce.Do(func() {
+		close(h.closeChan)
+	})
+}
+
+func (h *Handler) Wait() {
+	h.connectionsWg.Wait()
+}