golib/network/modbus/buffer.go

package modbus

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

	"golib/network"
)

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

type BuffHandler func(b []byte) error

type Buffer struct {
	Conn     net.Conn
	Handle   BuffHandler
	Cache    atomic.Value
	Creator  Creator
	Interval time.Duration
	Wait     chan []byte
	Logger   network.Logger

	stop    bool
	started bool
}

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", network.Bytes(b).HexTo())
	n, err := rw.Conn.Write(b)
	if err != nil {
		rw.Logger.Println("Write err: %s", err)
		return
	}
	if n != len(b) {
		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.Logger.Println("Read err: %s", err)
		return
	}
	rw.Cache.Store(body[:n])
	rw.Logger.Println("Read: %s", body[:n])
	if rw.Handle != nil {
		if err = rw.Handle(body[:n]); err != nil {
			rw.Logger.Println("TimerHandler 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) Stop() {
	rw.stop = true
}

func (rw *Buffer) Close() {
	rw.Stop()
	_ = rw.Conn.Close()
}

func (rw *Buffer) Start() {
	if rw.started {
		return
	}

	rw.callCreate() // call once

	if rw.Interval <= 0 {
		rw.Interval = network.WriteInterval
	}
	t := time.NewTimer(rw.Interval)
	for !rw.stop {
		select {
		case <-t.C:
			rw.callCreate()
			t.Reset(rw.Interval)
		case b := <-rw.Wait:
			rw.handleData(b)
		}
	}
	rw.started = false
}

func NewBuffer(conn net.Conn, creator Creator) *Buffer {
	buf := new(Buffer)
	buf.Conn = conn
	buf.Wait = make(chan []byte, 3)
	buf.Creator = creator
	buf.Logger = network.DefaultLogger
	return buf
}