proto-socket/go/test/heartbeat_test.go

package proto_socket_test

import (
	"context"
	"fmt"
	"io"
	"net"
	"net/http"
	"strings"
	"testing"
	"time"

	"nhooyr.io/websocket"

	toki "git.toki-labs.com/toki/proto-socket/go"
)

func TestHeartbeatDisconnectsWithoutResponse(t *testing.T) {
	clientConn, peerConn := net.Pipe()
	defer peerConn.Close()

	go func() {
		_, _ = io.Copy(io.Discard, peerConn)
	}()

	client := toki.NewTcpClient(clientConn, 1, 1, testParserMap())
	defer client.Close()

	disconnected := make(chan struct{}, 1)
	client.AddDisconnectListener(func(*toki.TcpClient) {
		disconnected <- struct{}{}
	})

	select {
	case <-disconnected:
	case <-time.After(3500 * time.Millisecond):
		t.Fatal("heartbeat timeout did not disconnect client")
	}
	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) {
	port := freePort(t)
	mux := http.NewServeMux()
	mux.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
		conn, err := websocket.Accept(w, r, nil)
		if err != nil {
			return
		}
		defer conn.Close(websocket.StatusNormalClosure, "")
		for {
			if _, _, err := conn.Read(r.Context()); err != nil {
				return
			}
		}
	})
	server := &http.Server{Addr: fmt.Sprintf("127.0.0.1:%d", port), Handler: mux}
	ln, err := net.Listen("tcp", server.Addr)
	if err != nil {
		t.Fatal(err)
	}
	defer server.Close()
	go func() {
		_ = server.Serve(ln)
	}()

	ctx, cancel := context.WithCancel(context.Background())
	defer cancel()
	client, err := toki.DialWsWithHeartbeat(ctx, "127.0.0.1", port, "/", 1, 1, testParserMap())
	if err != nil {
		t.Fatal(err)
	}
	defer client.Close()

	disconnected := make(chan struct{}, 1)
	client.AddDisconnectListener(func(*toki.WsClient) {
		disconnected <- struct{}{}
	})

	select {
	case <-disconnected:
	case <-time.After(3500 * time.Millisecond):
		t.Fatal("websocket heartbeat timeout did not disconnect client")
	}
	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)
	}
}