diff --git a/tcp/echo_test.go b/tcp/echo_test.go index 1bc211c8..c96ae6cb 100644 --- a/tcp/echo_test.go +++ b/tcp/echo_test.go @@ -51,3 +51,40 @@ func TestListenAndServe(t *testing.T) { closeChan <- struct{}{} time.Sleep(time.Second) } + +func TestClientCounter(t *testing.T) { + const addr = "127.0.0.1:8080" + go func() { + err := ListenAndServeWithSignal(&Config{ + Address: addr, + MaxConnect: 10000, + Timeout: 1 * time.Second, + }, MakeEchoHandler()) + if err != nil { + t.Errorf("listen and serve error: %s", err.Error()) + } + }() + + sleepUntil := time.Now().Add(3 * time.Second) + subtime := func() time.Duration { + return sleepUntil.Sub(time.Now()) + } + + for i := 0; i < 1000; i++ { + go func() { + conn, err := net.Dial("tcp", addr) + if err != nil { + t.Errorf(err.Error()) + } + defer conn.Close() + + time.Sleep(subtime()) + }() + time.Sleep(5 * time.Microsecond) + } + + time.Sleep(3 * time.Second) + if ClientCounter != 0 { + t.Errorf("client counter error: %d", ClientCounter) + } +} diff --git a/tcp/server.go b/tcp/server.go index e381de17..b74de25d 100644 --- a/tcp/server.go +++ b/tcp/server.go @@ -11,6 +11,7 @@ import ( "os" "os/signal" "sync" + "sync/atomic" "syscall" "time" @@ -26,7 +27,7 @@ type Config struct { } // ClientCounter Record the number of clients in the current Godis server -var ClientCounter int +var ClientCounter int32 // ListenAndServeWithSignal binds port and handle requests, blocking until receive stop signal func ListenAndServeWithSignal(cfg *Config, handler tcp.Handler) error { @@ -88,7 +89,7 @@ func ListenAndServe(listener net.Listener, handler tcp.Handler, closeChan <-chan go func() { defer func() { waitDone.Done() - ClientCounter-- + atomic.AddInt32(&ClientCounter, -1) }() handler.Handle(ctx, conn) }()