proto-socket/go/base_client.go

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
	heartbeatWait       time.Duration
	waitingHBResponse   bool
	hbTimer             *HeartbeatTimer
	hbMu                sync.Mutex
	disconnectListeners []func(Self)
	disconnectInfo      DisconnectInfo
	disconnectMu        sync.Mutex
	connCloseOnce       sync.Once
	self                Self
	doClose             func() error
}

func newBaseClient[Self any](self Self, intervalSec, waitSec int, doClose func() error) baseClient[Self] {
	return baseClient[Self]{
		heartbeatInterval: time.Duration(intervalSec) * time.Second,
		heartbeatWait:     time.Duration(waitSec) * time.Second,
		self:              self,
		doClose:           doClose,
	}
}

func (c *baseClient[Self]) AddDisconnectListener(handler func(Self)) {
	c.disconnectMu.Lock()
	defer c.disconnectMu.Unlock()
	c.disconnectListeners = append(c.disconnectListeners, handler)
}

func (c *baseClient[Self]) RemoveDisconnectListeners() {
	c.disconnectMu.Lock()
	defer c.disconnectMu.Unlock()
	c.disconnectListeners = nil
}

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()
		if c.doClose != nil {
			err = c.doClose()
		}
		c.notifyDisconnected()
	})
	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
	}
	c.hbMu.Lock()
	if c.hbTimer != nil {
		c.hbTimer.Stop()
	}
	c.hbTimer = NewHeartbeatTimer(c.heartbeatInterval, func() {
		if !c.IsAlive() {
			return
		}
		c.hbMu.Lock()
		c.waitingHBResponse = true
		c.hbMu.Unlock()
		_ = c.Send(&packets.HeartBeat{})
		c.hbMu.Lock()
		if c.hbTimer != nil {
			c.hbTimer.Stop()
		}
		c.hbTimer = NewHeartbeatTimer(c.heartbeatWait, func() {
			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()
	})
	c.hbMu.Unlock()
}

func (c *baseClient[Self]) onHeartBeat() {
	c.hbMu.Lock()
	if c.waitingHBResponse {
		c.waitingHBResponse = false
		c.hbMu.Unlock()
		return
	}
	c.hbMu.Unlock()
	_ = c.Send(&packets.HeartBeat{})
}

func (c *baseClient[Self]) stopHeartbeat() {
	c.hbMu.Lock()
	defer c.hbMu.Unlock()
	if c.hbTimer != nil {
		c.hbTimer.Stop()
		c.hbTimer = nil
	}
}

func (c *baseClient[Self]) onDisconnected(reason string, cause error) {
	_ = c.closeWithInfo(reason, cause)
}

func (c *baseClient[Self]) notifyDisconnected() {
	c.disconnectMu.Lock()
	listeners := append([]func(Self){}, c.disconnectListeners...)
	c.disconnectListeners = nil
	c.disconnectMu.Unlock()

	for _, listener := range listeners {
		listener(c.self)
	}
}