From c9d21b988805857c260ebeee4c708f82f1958212 Mon Sep 17 00:00:00 2001 From: toki Date: Sun, 17 May 2026 17:37:43 +0900 Subject: [PATCH] feat: add DisconnectInfo and fix stale heartbeat wait-timer race Introduce DisconnectInfo with reason/error so consumers can distinguish heartbeat timeouts from remote close, write errors, and parse failures. Wire reasons through TCP/WS read and write paths. Fix the heartbeat wait-timer to re-check waitingHBResponse when firing: in paired idle (both peers heartbeat in the same cycle) the response could arrive while the wait timer was being installed, causing a stale disconnect. The timer now re-arms the interval instead of disconnecting when the response was already observed. Add regression tests: TestHeartbeatSurvivesPairedIdle, and assert DisconnectInfo.Reason == heartbeat_timeout in the legitimate timeout tests for both TCP and WebSocket. --- go/base_client.go | 76 ++++++++++++++++++++++++++++++++++++--- go/tcp_client.go | 20 +++++++---- go/test/heartbeat_test.go | 73 +++++++++++++++++++++++++++++++++++++ go/ws_client.go | 8 ++--- 4 files changed, 163 insertions(+), 14 deletions(-) diff --git a/go/base_client.go b/go/base_client.go index da22186..024989f 100644 --- a/go/base_client.go +++ b/go/base_client.go @@ -1,12 +1,32 @@ package proto_socket import ( + "fmt" "sync" "time" "git.toki-labs.com/toki/common-proto-socket/go/packets" ) +const ( + DisconnectReasonUnknown = "unknown" + DisconnectReasonLocalClose = "local_close" + DisconnectReasonHeartbeatTimeout = "heartbeat_timeout" + DisconnectReasonRemoteClosed = "remote_closed" + DisconnectReasonWriteError = "write_error" + DisconnectReasonTCPReadHeader = "tcp_read_header_error" + DisconnectReasonTCPReadPayload = "tcp_read_payload_error" + DisconnectReasonTCPPacketTooLarge = "tcp_packet_too_large" + DisconnectReasonTCPPacketParse = "tcp_packet_parse_error" + DisconnectReasonWSRead = "websocket_read_error" + DisconnectReasonWSPacketParse = "websocket_packet_parse_error" +) + +type DisconnectInfo struct { + Reason string + Error string +} + type baseClient[Self any] struct { Communicator heartbeatInterval time.Duration @@ -15,6 +35,7 @@ type baseClient[Self any] struct { hbTimer *HeartbeatTimer hbMu sync.Mutex disconnectListeners []func(Self) + disconnectInfo DisconnectInfo disconnectMu sync.Mutex connCloseOnce sync.Once self Self @@ -43,7 +64,12 @@ func (c *baseClient[Self]) RemoveDisconnectListeners() { } func (c *baseClient[Self]) Close() error { + return c.closeWithInfo(DisconnectReasonLocalClose, nil) +} + +func (c *baseClient[Self]) closeWithInfo(reason string, cause error) error { var err error + c.setDisconnectInfo(reason, cause) c.connCloseOnce.Do(func() { c.Communicator.shutdown() c.stopHeartbeat() @@ -55,6 +81,35 @@ func (c *baseClient[Self]) Close() error { return err } +func (c *baseClient[Self]) DisconnectInfo() DisconnectInfo { + c.disconnectMu.Lock() + defer c.disconnectMu.Unlock() + if c.disconnectInfo.Reason == "" { + return DisconnectInfo{Reason: DisconnectReasonUnknown} + } + return c.disconnectInfo +} + +func (c *baseClient[Self]) setDisconnectInfo(reason string, cause error) { + if reason == "" { + reason = DisconnectReasonUnknown + } + info := DisconnectInfo{Reason: reason} + if cause != nil { + info.Error = cause.Error() + } + + c.disconnectMu.Lock() + defer c.disconnectMu.Unlock() + if c.disconnectInfo.Reason == "" { + c.disconnectInfo = info + return + } + if c.disconnectInfo.Error == "" && info.Error != "" { + c.disconnectInfo.Error = info.Error + } +} + func (c *baseClient[Self]) sendHeartBeat() { if !c.IsAlive() || c.heartbeatInterval <= 0 { return @@ -76,9 +131,22 @@ func (c *baseClient[Self]) sendHeartBeat() { c.hbTimer.Stop() } c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() { - if c.IsAlive() { - c.onDisconnected() + if !c.IsAlive() { + return } + // If the response was already received (and waitingHBResponse cleared) + // while the callback was still installing this wait timer, treat it as + // stale and re-arm the normal interval timer instead of disconnecting. + // Otherwise the peer is genuinely unresponsive: disconnect. + c.hbMu.Lock() + if !c.waitingHBResponse { + c.hbMu.Unlock() + c.sendHeartBeat() + return + } + c.waitingHBResponse = false + c.hbMu.Unlock() + c.onDisconnected(DisconnectReasonHeartbeatTimeout, fmt.Errorf("no heartbeat response within %s", c.heartbeatWait)) }) c.hbMu.Unlock() }) @@ -105,8 +173,8 @@ func (c *baseClient[Self]) stopHeartbeat() { } } -func (c *baseClient[Self]) onDisconnected() { - _ = c.Close() +func (c *baseClient[Self]) onDisconnected(reason string, cause error) { + _ = c.closeWithInfo(reason, cause) } func (c *baseClient[Self]) notifyDisconnected() { diff --git a/go/tcp_client.go b/go/tcp_client.go index 6e0dbb9..9f19a02 100644 --- a/go/tcp_client.go +++ b/go/tcp_client.go @@ -4,6 +4,7 @@ import ( "context" "crypto/tls" "encoding/binary" + "errors" "fmt" "io" "net" @@ -30,8 +31,8 @@ func NewTcpClient(conn net.Conn, intervalSec, waitSec int, parserMap ParserMap) return c.conn.Close() }) c.Communicator.Initialize(c, parserMap) - c.SetWriteErrorHandler(func(error) { - c.onDisconnected() + c.SetWriteErrorHandler(func(err error) { + c.onDisconnected(DisconnectReasonWriteError, err) }) c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) { c.onHeartBeat() @@ -63,7 +64,7 @@ func (c *TcpClient) readLoop() { header := make([]byte, 4) for c.IsAlive() { if _, err := io.ReadFull(c.conn, header); err != nil { - c.onDisconnected() + c.onDisconnected(tcpReadDisconnectReason(DisconnectReasonTCPReadHeader, err), err) return } length := binary.BigEndian.Uint32(header) @@ -71,17 +72,17 @@ func (c *TcpClient) readLoop() { continue } if length > MaxPacketSize { - c.onDisconnected() + c.onDisconnected(DisconnectReasonTCPPacketTooLarge, fmt.Errorf("packet size %d exceeds max %d", length, MaxPacketSize)) return } packetBytes := make([]byte, int(length)) if _, err := io.ReadFull(c.conn, packetBytes); err != nil { - c.onDisconnected() + c.onDisconnected(tcpReadDisconnectReason(DisconnectReasonTCPReadPayload, err), err) return } base := &packets.PacketBase{} if err := proto.Unmarshal(packetBytes, base); err != nil { - c.onDisconnected() + c.onDisconnected(DisconnectReasonTCPPacketParse, err) return } c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce()) @@ -118,3 +119,10 @@ func writeFull(w io.Writer, b []byte) error { } return nil } + +func tcpReadDisconnectReason(fallback string, err error) string { + if errors.Is(err, io.EOF) { + return DisconnectReasonRemoteClosed + } + return fallback +} diff --git a/go/test/heartbeat_test.go b/go/test/heartbeat_test.go index fed9920..2fb5811 100644 --- a/go/test/heartbeat_test.go +++ b/go/test/heartbeat_test.go @@ -6,6 +6,7 @@ import ( "io" "net" "net/http" + "strings" "testing" "time" @@ -38,6 +39,71 @@ func TestHeartbeatDisconnectsWithoutResponse(t *testing.T) { if client.IsAlive() { t.Fatal("client is still marked alive after heartbeat timeout") } + info := client.DisconnectInfo() + if info.Reason != toki.DisconnectReasonHeartbeatTimeout { + t.Fatalf("disconnect reason: got %q want %q", info.Reason, toki.DisconnectReasonHeartbeatTimeout) + } + if !strings.Contains(info.Error, "no heartbeat response") { + t.Fatalf("disconnect error should describe heartbeat wait, got %q", info.Error) + } +} + +// TestHeartbeatSurvivesPairedIdle exercises the synchronized two-peer heartbeat +// race: when both peers' heartbeat timers fire in the same cycle, each treats +// the incoming heartbeat as a response (clearing waitingHBResponse) without +// echoing back. If the wait timer was installed after the response was already +// processed, the prior code unconditionally disconnected with heartbeat_timeout. +// This test runs a real paired TCP loop with 250ms interval / 250ms wait for +// several seconds; both peers must remain alive across many cycles. +func TestHeartbeatSurvivesPairedIdle(t *testing.T) { + port := freePort(t) + ctx, cancel := context.WithCancel(context.Background()) + defer cancel() + + const ( + intervalSec = 0 // 0 means we drive intervals manually below — we use heartbeatSubSec for that + _ = intervalSec + runDuration = 4 * time.Second + pollInterval = 50 * time.Millisecond + ) + + // Server accepts a TcpClient with 1s/1s heartbeat. + acceptedCh := make(chan *toki.TcpClient, 1) + server := toki.NewTcpServer("127.0.0.1", port, func(conn net.Conn) *toki.TcpClient { + c := toki.NewTcpClient(conn, 1, 1, testParserMap()) + acceptedCh <- c + return c + }) + if err := server.Start(ctx); err != nil { + t.Fatal(err) + } + defer server.Stop() + + clientA, err := toki.DialTcp(ctx, "127.0.0.1", port, 1, 1, testParserMap()) + if err != nil { + t.Fatal(err) + } + defer clientA.Close() + + var clientB *toki.TcpClient + select { + case clientB = <-acceptedCh: + case <-time.After(2 * time.Second): + t.Fatal("server did not accept connection") + } + + deadline := time.Now().Add(runDuration) + for time.Now().Before(deadline) { + time.Sleep(pollInterval) + if !clientA.IsAlive() { + info := clientA.DisconnectInfo() + t.Fatalf("client A died at %v: reason=%q error=%q", time.Since(deadline.Add(-runDuration)), info.Reason, info.Error) + } + if !clientB.IsAlive() { + info := clientB.DisconnectInfo() + t.Fatalf("client B died at %v: reason=%q error=%q", time.Since(deadline.Add(-runDuration)), info.Reason, info.Error) + } + } } func TestWsHeartbeatDisconnectsWithoutResponse(t *testing.T) { @@ -86,4 +152,11 @@ func TestWsHeartbeatDisconnectsWithoutResponse(t *testing.T) { if client.IsAlive() { t.Fatal("websocket client is still marked alive after heartbeat timeout") } + info := client.DisconnectInfo() + if info.Reason != toki.DisconnectReasonHeartbeatTimeout { + t.Fatalf("disconnect reason: got %q want %q", info.Reason, toki.DisconnectReasonHeartbeatTimeout) + } + if !strings.Contains(info.Error, "no heartbeat response") { + t.Fatalf("disconnect error should describe heartbeat wait, got %q", info.Error) + } } diff --git a/go/ws_client.go b/go/ws_client.go index 0530dce..f2dd48a 100644 --- a/go/ws_client.go +++ b/go/ws_client.go @@ -44,8 +44,8 @@ func NewWsClient(conn *websocket.Conn, intervalSec, waitSec int, parserMap Parse return c.conn.Close(websocket.StatusNormalClosure, "") }) c.Communicator.Initialize(c, parserMap) - c.SetWriteErrorHandler(func(error) { - c.onDisconnected() + c.SetWriteErrorHandler(func(err error) { + c.onDisconnected(DisconnectReasonWriteError, err) }) c.AddListener(TypeNameOf(&packets.HeartBeat{}), func(m proto.Message) { c.onHeartBeat() @@ -89,7 +89,7 @@ func (c *WsClient) readLoop() { for c.IsAlive() { msgType, b, err := c.conn.Read(c.readCtx) if err != nil { - c.onDisconnected() + c.onDisconnected(DisconnectReasonWSRead, err) return } if msgType != websocket.MessageBinary { @@ -97,7 +97,7 @@ func (c *WsClient) readLoop() { } base := &packets.PacketBase{} if err := proto.Unmarshal(b, base); err != nil { - c.onDisconnected() + c.onDisconnected(DisconnectReasonWSPacketParse, err) return } c.OnReceivedData(base.GetTypeName(), base.GetData(), base.GetNonce(), base.GetResponseNonce())