proto-socket/go/base_client.go

package proto_socket

import (
	"sync"
	"time"

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

type baseClient[Self any] struct {
	Communicator
	heartbeatInterval   time.Duration
	heartbeatWait       time.Duration
	waitingHBResponse   bool
	hbTimer             *HeartbeatTimer
	hbMu                sync.Mutex
	disconnectListeners []func(Self)
	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 {
	var err error
	c.connCloseOnce.Do(func() {
		c.Communicator.shutdown()
		c.stopHeartbeat()
		if c.doClose != nil {
			err = c.doClose()
		}
		c.notifyDisconnected()
	})
	return err
}

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() {
				c.onDisconnected()
			}
		})
		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() {
	_ = c.Close()
}

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)
	}
}