golib/gnet/modbus/buffer.go

package modbus

import (
	"context"
	"net"
	"sync/atomic"
	"time"

	"golib/gnet"
)

// Creator 创建需要写入的数据
type Creator interface {
	Create() ([]byte, error)
}

type BufHandler func(b []byte) error
type ErrHandler func(err error)

func defaultHandler(_ []byte) error { return nil }
func defaultErrHandle(_ error)      {}

type Buffer struct {
	Conn      net.Conn
	Handle    BufHandler // 读取数据后执行
	ErrHandle ErrHandler // 读写失败时执行
	Cache     atomic.Value
	Creator   Creator       // 当 Wait 无数据且到达轮询时间时执行
	Interval  time.Duration // 轮询频率
	Wait      chan []byte
	Logger    gnet.Logger

	ctx context.Context
}

func (rw *Buffer) Get() ([]byte, bool) {
	b, ok := rw.Cache.Load().([]byte)
	if !ok {
		return nil, false
	}
	return b, true
}

func (rw *Buffer) Send(b []byte) {
	rw.Wait <- b
}

func (rw *Buffer) handleData(b []byte) {
	rw.Logger.Println("Write: %s", gnet.Bytes(b).HexTo())

	n, err := rw.Conn.Write(b)
	if err != nil {
		rw.ErrHandle(err)
		rw.Logger.Println("Write err: %s", err)
		return
	}

	if n != len(b) {
		rw.ErrHandle(err)
		rw.Logger.Println("Write err: not fully write: data length: %d write length: %d", len(b), n)
		return
	}

	body := make([]byte, 4096)

	n, err = rw.Conn.Read(body)
	if err != nil {
		rw.ErrHandle(err)
		rw.Logger.Println("Read err: %s", err)
		return
	}

	rw.Cache.Store(body[:n])
	rw.Logger.Println("Read: %s", gnet.Bytes(body[:n]).HexTo())

	if err = rw.Handle(body[:n]); err != nil {
		rw.Logger.Println("Handle err: %s", err)
	}
}

func (rw *Buffer) callCreate() {
	if rw.Creator != nil {
		b, err := rw.Creator.Create()
		if err != nil {
			rw.Logger.Println("Handle Create err: %s", err)
		} else {
			rw.handleData(b)
		}
	}
}

func (rw *Buffer) Start() {
	rw.callCreate() // call once

	if rw.Interval <= 0 {
		rw.Interval = gnet.WriteInterval
	}

	t := time.NewTimer(rw.Interval)
	defer t.Stop()

	for {
		select {
		case <-rw.ctx.Done():
			_ = rw.Conn.Close()
			rw.ErrHandle(rw.ctx.Err())
			return
		case <-t.C:
			rw.callCreate()
			t.Reset(rw.Interval)
		case b := <-rw.Wait:
			rw.handleData(b)
		}
	}
}

func NewBuffer(ctx context.Context, conn net.Conn, creator Creator) *Buffer {
	buf := new(Buffer)
	buf.Conn = conn
	buf.Handle = defaultHandler
	buf.ErrHandle = defaultErrHandle
	buf.Wait = make(chan []byte, 3)
	buf.Creator = creator
	buf.Logger = gnet.DefaultLogger("[Buffer] ")
	buf.ctx = ctx
	return buf
}