nnet/middleware/heartbeat.go

package middleware

import (
	"git.noahlan.cn/noahlan/nnet/core"
	"git.noahlan.cn/noahlan/nnet/entity"
	"git.noahlan.cn/noahlan/nnet/packet"
	"git.noahlan.cn/noahlan/ntools-go/core/nlog"
	"sync/atomic"
	"time"
)

type HeartbeatMiddleware struct {
	lastAt   int64
	interval time.Duration
	hbdFn    func(entity entity.NetworkEntity) []byte
}

func WithHeartbeat(interval time.Duration, hbdFn func(entity entity.NetworkEntity) []byte) core.RunOption {
	m := &HeartbeatMiddleware{
		lastAt:   time.Now().Unix(),
		interval: interval,
		hbdFn:    hbdFn,
	}
	if hbdFn == nil {
		nlog.Error("dataFn must not be nil")
		panic("dataFn must not be nil")
	}
	core.Lifetime.OnOpen(m.start)

	return func(server *core.NNet) {
		server.Use(func(next core.HandlerFunc) core.HandlerFunc {
			return func(entity entity.NetworkEntity, pkg packet.IPacket) {
				m.handle(entity, pkg)

				next(entity, pkg)
			}
		})
	}
}

func (m *HeartbeatMiddleware) start(entity entity.NetworkEntity) {
	ticker := time.NewTicker(m.interval)

	defer func() {
		ticker.Stop()
	}()

	for {
		select {
		case <-ticker.C:
			deadline := time.Now().Add(-2 * m.interval).Unix()
			if atomic.LoadInt64(&m.lastAt) < deadline {
				nlog.Errorf("Heartbeat timeout, LastTime=%d, Deadline=%d", atomic.LoadInt64(&m.lastAt), deadline)
				return
			}
			err := entity.SendBytes(m.hbdFn(entity))
			if err != nil {
				nlog.Errorf("Heartbeat err: %v", err)
				return
			}
		}
	}
}

func (m *HeartbeatMiddleware) handle(_ entity.NetworkEntity, _ packet.IPacket) {
	atomic.StoreInt64(&m.lastAt, time.Now().Unix())
}