diff --git a/client.go b/client.go index 37836e09..d52f09f0 100644 --- a/client.go +++ b/client.go @@ -22,7 +22,6 @@ import ( "crypto/x509" "encoding/pem" "fmt" - "math" "net" "os" "strings" @@ -153,29 +152,32 @@ func (c *client) dialTCP() Session { conn net.Conn ) - for { - if c.IsClosed() { - return nil - } - if c.sslEnabled { - if sslConfig, buildTlsConfErr := c.tlsConfigBuilder.BuildTlsConfig(); buildTlsConfErr == nil && sslConfig != nil { - d := &net.Dialer{Timeout: connectTimeout} - conn, err = tls.DialWithDialer(d, "tcp", c.addr, sslConfig) - } + if c.IsClosed() { + return nil + } + if c.sslEnabled { + sslConfig, buildTlsConfErr := c.tlsConfigBuilder.BuildTlsConfig() + if buildTlsConfErr != nil { + err = buildTlsConfErr + } else if sslConfig == nil { + err = fmt.Errorf("tlsConfigBuilder returned nil config without error") } else { - conn, err = net.DialTimeout("tcp", c.addr, connectTimeout) - } - if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { - _ = conn.Close() - err = errSelfConnect + d := &net.Dialer{Timeout: connectTimeout} + conn, err = tls.DialWithDialer(d, "tcp", c.addr, sslConfig) } - if err == nil { - return newTCPSession(conn, c) - } - - log.Infof("net.DialTimeout(addr:%s, timeout:%v) = error:%+v", c.addr, connectTimeout, perrors.WithStack(err)) - <-gxtime.After(connectInterval) + } else { + conn, err = net.DialTimeout("tcp", c.addr, connectTimeout) + } + if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { + _ = conn.Close() + err = errSelfConnect + } + if err == nil { + return newTCPSession(conn, c) } + + log.Infof("net.DialTimeout(addr:%s, timeout:%v) = error:%+v", c.addr, connectTimeout, perrors.WithStack(err)) + return nil } func (c *client) dialUDP() Session { @@ -194,46 +196,41 @@ func (c *client) dialUDP() Session { buf = *bufp localAddr = &net.UDPAddr{IP: net.IPv4zero, Port: 0} peerAddr, _ = net.ResolveUDPAddr("udp", c.addr) - for { - if c.IsClosed() { - return nil - } - conn, err = net.DialUDP("udp", localAddr, peerAddr) - if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { - _ = conn.Close() - err = errSelfConnect - } - if err != nil { - log.Warnf("net.DialTimeout(addr:%s, timeout:%v) = error:%+v", c.addr, perrors.WithStack(err)) - <-gxtime.After(connectInterval) - continue - } + if c.IsClosed() { + return nil + } + conn, err = net.DialUDP("udp", localAddr, peerAddr) + if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { + _ = conn.Close() + err = errSelfConnect + } + if err != nil { + log.Warnf("net.DialUDP(addr:%s) = error:%+v", c.addr, perrors.WithStack(err)) + return nil + } - // check connection alive by write/read action - if err := conn.SetWriteDeadline(time.Now().Add(1e9)); err != nil { - log.Warnf("failed to set write deadline: %+v", err) - } - if length, err = conn.Write(connectPingPackage[:]); err != nil { - _ = conn.Close() - log.Warnf("conn.Write(%s) = {length:%d, err:%+v}", string(connectPingPackage), length, perrors.WithStack(err)) - <-gxtime.After(connectInterval) - continue - } - if err := conn.SetReadDeadline(time.Now().Add(1e9)); err != nil { - log.Warnf("failed to set read deadline: %+v", err) - } - length, err = conn.Read(buf) - if netErr, ok := perrors.Cause(err).(net.Error); ok && netErr.Timeout() { - err = nil - } - if err != nil { - log.Infof("conn{%#v}.Read() = {length:%d, err:%+v}", conn, length, perrors.WithStack(err)) - _ = conn.Close() - <-gxtime.After(connectInterval) - continue - } - return newUDPSession(conn, c) + // check connection alive by write/read action + if err := conn.SetWriteDeadline(time.Now().Add(1e9)); err != nil { + log.Warnf("failed to set write deadline: %+v", err) + } + if length, err = conn.Write(connectPingPackage[:]); err != nil { + _ = conn.Close() + log.Warnf("conn.Write(%s) = {length:%d, err:%+v}", string(connectPingPackage), length, perrors.WithStack(err)) + return nil } + if err := conn.SetReadDeadline(time.Now().Add(1e9)); err != nil { + log.Warnf("failed to set read deadline: %+v", err) + } + length, err = conn.Read(buf) + if netErr, ok := perrors.Cause(err).(net.Error); ok && netErr.Timeout() { + err = nil + } + if err != nil { + log.Infof("conn{%#v}.Read() = {length:%d, err:%+v}", conn, length, perrors.WithStack(err)) + _ = conn.Close() + return nil + } + return newUDPSession(conn, c) } func (c *client) dialWS() Session { @@ -245,28 +242,25 @@ func (c *client) dialWS() Session { ) dialer.EnableCompression = true - for { - if c.IsClosed() { - return nil - } - conn, _, err = dialer.Dial(c.addr, nil) - log.Infof("websocket.dialer.Dial(addr:%s) = error:%+v", c.addr, perrors.WithStack(err)) - if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { - _ = conn.Close() - err = errSelfConnect - } - if err == nil { - ss = newWSSession(conn, c) - if ss.(*session).maxMsgLen > 0 { - conn.SetReadLimit(int64(ss.(*session).maxMsgLen)) - } - - return ss + if c.IsClosed() { + return nil + } + conn, _, err = dialer.Dial(c.addr, nil) + if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { + _ = conn.Close() + err = errSelfConnect + } + if err == nil { + ss = newWSSession(conn, c) + if ss.(*session).maxMsgLen > 0 { + conn.SetReadLimit(int64(ss.(*session).maxMsgLen)) } - log.Infof("websocket.dialer.Dial(addr:%s) = error:%+v", c.addr, perrors.WithStack(err)) - <-gxtime.After(connectInterval) + return ss } + + log.Infof("websocket.dialer.Dial(addr:%s) = error:%+v", c.addr, perrors.WithStack(err)) + return nil } func (c *client) dialWSS() Session { @@ -323,28 +317,26 @@ func (c *client) dialWSS() Session { // dialer.EnableCompression = true dialer.TLSClientConfig = config - for { - if c.IsClosed() { - return nil - } - conn, _, err = dialer.Dial(c.addr, nil) - if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { - _ = conn.Close() - err = errSelfConnect - } - if err == nil { - ss = newWSSession(conn, c) - if ss.(*session).maxMsgLen > 0 { - conn.SetReadLimit(int64(ss.(*session).maxMsgLen)) - } - ss.SetName(defaultWSSSessionName) - - return ss + if c.IsClosed() { + return nil + } + conn, _, err = dialer.Dial(c.addr, nil) + if err == nil && gxnet.IsSameAddr(conn.RemoteAddr(), conn.LocalAddr()) { + _ = conn.Close() + err = errSelfConnect + } + if err == nil { + ss = newWSSession(conn, c) + if ss.(*session).maxMsgLen > 0 { + conn.SetReadLimit(int64(ss.(*session).maxMsgLen)) } + ss.SetName(defaultWSSSessionName) - log.Infof("websocket.dialer.Dial(addr:%s) = error:%+v", c.addr, perrors.WithStack(err)) - <-gxtime.After(connectInterval) + return ss } + + log.Infof("websocket.dialer.Dial(addr:%s) = error:%+v", c.addr, perrors.WithStack(err)) + return nil } func (c *client) dial() Session { @@ -381,36 +373,36 @@ func (c *client) sessionNum() int { return num } -func (c *client) connect() { +func (c *client) connect() bool { var ( err error ss Session ) - for { - ss = c.dial() - if ss == nil { - // client has been closed - break - } - err = c.newSession(ss) - if err == nil { - ss.(*session).run() - c.Lock() - if c.ssMap == nil { - c.Unlock() - break - } - c.ssMap[ss] = struct{}{} + ss = c.dial() + if ss == nil { + // client has been closed + return false + } + err = c.newSession(ss) + if err == nil { + ss.(*session).run() + c.Lock() + if c.ssMap == nil { c.Unlock() - ss.SetAttribute(sessionClientKey, c) - ss.SetAttribute(ignoreReconnectKey, false) - break + ss.Close() + return false } - // don't distinguish between tcp connection and websocket connection. Because - // gorilla/websocket/conn.go:(Conn)Close also invoke net.Conn.Close() - _ = ss.Conn().Close() + c.ssMap[ss] = struct{}{} + c.Unlock() + ss.SetAttribute(sessionClientKey, c) + ss.SetAttribute(ignoreReconnectKey, false) + return true } + // don't distinguish between tcp connection and websocket connection. Because + // gorilla/websocket/conn.go:(Conn)Close also invoke net.Conn.Close() + _ = ss.Conn().Close() + return false } // there are two methods to keep connection pool. the first approach is like @@ -429,10 +421,7 @@ func (c *client) RunEventLoop(newSession NewSessionCallback) { // a for-loop connect to make sure the connection pool is valid func (c *client) reConnect() { - var ( - sessionNum, reconnectAttempts int - maxReconnectInterval int64 - ) + var reconnectAttempts int reconnectInterval := c.reconnectInterval if reconnectInterval == 0 { reconnectInterval = defaultReconnectInterval @@ -442,21 +431,34 @@ func (c *client) reConnect() { maxReconnectAttempts = defaultMaxReconnectAttempts } connPoolSize := c.number - for { + for reconnectAttempts < maxReconnectAttempts { if c.IsClosed() { log.Warnf("client{peer:%s} goroutine exit now.", c.addr) - break + return } - sessionNum = c.sessionNum() - if connPoolSize <= sessionNum || maxReconnectAttempts < reconnectAttempts { - //exit reconnect when the number of connection pools is sufficient or the current reconnection attempts exceeds the max reconnection attempts. - break + if connPoolSize <= c.sessionNum() { + return + } + if c.connect() { + reconnectAttempts = 0 + continue } - c.connect() reconnectAttempts++ - maxReconnectInterval = int64(math.Min(float64(reconnectAttempts), float64(maxBackOffTimes))) * int64(reconnectInterval) - <-gxtime.After(time.Duration(maxReconnectInterval)) + if c.IsClosed() || connPoolSize <= c.sessionNum() || reconnectAttempts >= maxReconnectAttempts { + return + } + + backOffTimes := reconnectAttempts + if maxBackOffTimes < backOffTimes { + backOffTimes = maxBackOffTimes + } + backoff := time.Duration(int64(backOffTimes) * int64(reconnectInterval)) + select { + case <-c.done: + return + case <-gxtime.After(backoff): + } } } diff --git a/client_test.go b/client_test.go index a0a3336d..2eaebba0 100644 --- a/client_test.go +++ b/client_test.go @@ -19,6 +19,7 @@ package getty import ( "bytes" + "errors" "net" "net/http" "os" @@ -69,6 +70,21 @@ func (h *MessageHandler) OnClose(session Session) {} func (h *MessageHandler) OnMessage(session Session, pkg any) {} func (h *MessageHandler) OnCron(session Session) {} +type poolClosingMessageHandler struct { + MessageHandler + client *client + opened chan Session +} + +func (h *poolClosingMessageHandler) OnOpen(session Session) error { + h.client.Lock() + h.client.ssMap = nil + h.client.Unlock() + h.opened <- session + + return h.MessageHandler.OnOpen(session) +} + type Package struct{} func (p Package) String() string { @@ -93,7 +109,7 @@ func newSessionCallback(session Session, handler *MessageHandler) error { func TestTCPClient(t *testing.T) { listenLocalServer := func() (net.Listener, error) { - listener, err := net.Listen("tcp", ":0") + listener, err := net.Listen("tcp", "127.0.0.1:0") if err != nil { return nil, err } @@ -201,6 +217,271 @@ func TestTCPClient(t *testing.T) { assert.True(t, clt.IsClosed()) } +func TestTCPClientReconnectStopsAfterMaxAttempts(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + assert.Nil(t, err) + assert.NotNil(t, listener) + addr := listener.Addr().String() + assert.Nil(t, listener.Close()) + + clt := NewTCPClient( + WithServerAddress(addr), + WithReconnectInterval(int(time.Millisecond)), + WithConnectionNumber(1), + WithReconnectAttempts(2), + ) + assert.NotNil(t, clt) + + done := make(chan struct{}) + go func() { + var msgHandler MessageHandler + cb := func(session Session) error { + return newSessionCallback(session, &msgHandler) + } + clt.RunEventLoop(cb) + close(done) + }() + + select { + case <-done: + case <-time.After(2 * time.Second): + t.Fatal("client reconnect loop did not stop after max reconnect attempts") + } + + clt.Close() + assert.True(t, clt.IsClosed()) +} + +func TestTCPClientReconnectBacksOffWhenNewSessionFails(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + assert.Nil(t, err) + assert.NotNil(t, listener) + + acceptDone := make(chan struct{}) + accepted := make(chan struct{}, 2) + go func() { + defer close(acceptDone) + for { + conn, err := listener.Accept() + if err != nil { + return + } + select { + case accepted <- struct{}{}: + default: + } + _ = conn.Close() + } + }() + defer func() { + _ = listener.Close() + <-acceptDone + }() + + clt := NewTCPClient( + WithServerAddress(listener.Addr().String()), + WithReconnectInterval(int(time.Millisecond)), + WithConnectionNumber(1), + WithReconnectAttempts(2), + ) + assert.NotNil(t, clt) + + done := make(chan struct{}) + attempts := 0 + go func() { + cb := func(session Session) error { + attempts++ + return errors.New("test new session failure") + } + clt.RunEventLoop(cb) + close(done) + }() + + select { + case <-done: + case <-time.After(500 * time.Millisecond): + t.Fatal("client reconnect loop did not stop when newSession kept failing") + } + + for i := 0; i < 2; i++ { + select { + case <-accepted: + case <-time.After(time.Second): + t.Fatal("server did not accept client connection") + } + } + assert.Equal(t, 2, attempts) + clt.Close() + assert.True(t, clt.IsClosed()) +} + +func TestTCPClientConnectClosesSessionWhenPoolClosed(t *testing.T) { + listener, err := net.Listen("tcp", "127.0.0.1:0") + assert.Nil(t, err) + assert.NotNil(t, listener) + defer func() { + _ = listener.Close() + }() + + acceptedConn := make(chan net.Conn, 1) + acceptErr := make(chan error, 1) + go func() { + conn, err := listener.Accept() + if err != nil { + acceptErr <- err + return + } + acceptedConn <- conn + }() + + clt := NewTCPClient( + WithServerAddress(listener.Addr().String()), + WithConnectionNumber(1), + ).(*client) + defer clt.Close() + + handler := &poolClosingMessageHandler{ + client: clt, + opened: make(chan Session, 1), + } + clt.newSession = func(session Session) error { + var pkgHandler PackageHandler + session.SetName("pool-closing-client-session") + session.SetPkgHandler(&pkgHandler) + session.SetEventListener(handler) + session.SetReadTimeout(time.Millisecond) + session.SetWriteTimeout(time.Millisecond) + session.SetWaitTime(time.Millisecond) + + return nil + } + + assert.False(t, clt.connect()) + + select { + case session := <-handler.opened: + assert.True(t, session.IsClosed()) + case <-time.After(time.Second): + t.Fatal("session was not opened") + } + + select { + case conn := <-acceptedConn: + _ = conn.Close() + case err := <-acceptErr: + t.Fatalf("server accept failed: %+v", err) + case <-time.After(time.Second): + t.Fatal("server did not accept client connection") + } +} + +func TestTCPClientReconnectAttemptsIgnoreSuccessfulPoolFills(t *testing.T) { + const connPoolSize = 3 + listener, err := net.Listen("tcp", "127.0.0.1:0") + assert.Nil(t, err) + assert.NotNil(t, listener) + addr := listener.Addr().String() + assert.Nil(t, listener.Close()) + + var ( + serverLock sync.Mutex + listeners []net.Listener + conns []net.Conn + acceptDone <-chan struct{} + serverErr = make(chan error, connPoolSize) + ) + defer func() { + serverLock.Lock() + defer serverLock.Unlock() + for _, ln := range listeners { + _ = ln.Close() + } + for _, conn := range conns { + _ = conn.Close() + } + }() + + startServer := func() { + ln, err := net.Listen("tcp", addr) + if err != nil { + serverErr <- err + return + } + + done := make(chan struct{}) + serverLock.Lock() + listeners = append(listeners, ln) + acceptDone = done + serverLock.Unlock() + + go func() { + conn, err := ln.Accept() + _ = ln.Close() + if err == nil { + serverLock.Lock() + conns = append(conns, conn) + serverLock.Unlock() + } + close(done) + }() + } + + waitAccepted := func() { + serverLock.Lock() + done := acceptDone + acceptDone = nil + serverLock.Unlock() + if done == nil { + t.Fatal("client connected before the test server was ready") + } + select { + case <-done: + case <-time.After(time.Second): + t.Fatal("server did not accept client connection") + } + } + + time.AfterFunc(10*time.Millisecond, startServer) + + clt := NewTCPClient( + WithServerAddress(addr), + WithReconnectInterval(int(50*time.Millisecond)), + WithConnectionNumber(connPoolSize), + WithReconnectAttempts(2), + ) + assert.NotNil(t, clt) + defer clt.Close() + + var ( + msgHandler MessageHandler + successes int + ) + cb := func(session Session) error { + waitAccepted() + successes++ + if successes < connPoolSize { + time.AfterFunc(10*time.Millisecond, startServer) + } + return newSessionCallback(session, &msgHandler) + } + + done := make(chan struct{}) + go func() { + clt.RunEventLoop(cb) + close(done) + }() + + select { + case <-done: + case err := <-serverErr: + t.Fatalf("failed to start test server: %+v", err) + case <-time.After(3 * time.Second): + t.Fatal("client reconnect loop did not complete") + } + + assert.Equal(t, connPoolSize, msgHandler.SessionNumber()) +} + func TestUDPClient(t *testing.T) { var ( err error