starnet/que.go

package starnet

import (
	"bytes"
	"context"
	"encoding/binary"
	"errors"
	"os"
	"sync"
	"time"
)

// 识别头
var header = []byte{11, 27, 19, 96, 12, 25, 02, 20}

// MsgQueue 为基本的信息单位
type MsgQueue struct {
	ID   uint16
	Msg  []byte
	Conn interface{}
}

// StarQueue 为流数据中的消息队列分发
type StarQueue struct {
	count      int64
	Encode     bool
	Reserve    uint16
	Msgid      uint16
	MsgPool    chan MsgQueue
	UnFinMsg   sync.Map
	LastID     int //= -1
	ctx        context.Context
	cancel     context.CancelFunc
	duration   time.Duration
	EncodeFunc func([]byte) []byte
	DecodeFunc func([]byte) []byte
	//restoreMu  sync.Mutex
}

func NewQueueCtx(ctx context.Context, count int64) *StarQueue {
	var que StarQueue
	que.Encode = false
	que.count = count
	que.MsgPool = make(chan MsgQueue, count)
	if ctx == nil {
		que.ctx, que.cancel = context.WithCancel(context.Background())
	} else {
		que.ctx, que.cancel = context.WithCancel(ctx)
	}
	que.duration = 0
	return &que
}
func NewQueueWithCount(count int64) *StarQueue {
	return NewQueueCtx(nil, count)
}

// NewQueue 建立一个新消息队列
func NewQueue() *StarQueue {
	return NewQueueWithCount(32)
}

// Uint32ToByte 4位uint32转byte
func Uint32ToByte(src uint32) []byte {
	res := make([]byte, 4)
	res[3] = uint8(src)
	res[2] = uint8(src >> 8)
	res[1] = uint8(src >> 16)
	res[0] = uint8(src >> 24)
	return res
}

// ByteToUint32 byte转4位uint32
func ByteToUint32(src []byte) uint32 {
	var res uint32
	buffer := bytes.NewBuffer(src)
	binary.Read(buffer, binary.BigEndian, &res)
	return res
}

// Uint16ToByte 2位uint16转byte
func Uint16ToByte(src uint16) []byte {
	res := make([]byte, 2)
	res[1] = uint8(src)
	res[0] = uint8(src >> 8)
	return res
}

// ByteToUint16 用于byte转uint16
func ByteToUint16(src []byte) uint16 {
	var res uint16
	buffer := bytes.NewBuffer(src)
	binary.Read(buffer, binary.BigEndian, &res)
	return res
}

// BuildMessage 生成编码后的信息用于发送
func (que *StarQueue) BuildMessage(src []byte) []byte {
	var buff bytes.Buffer
	que.Msgid++
	if que.Encode {
		src = que.EncodeFunc(src)
	}
	length := uint32(len(src))
	buff.Write(header)
	buff.Write(Uint32ToByte(length))
	buff.Write(Uint16ToByte(que.Msgid))
	buff.Write(src)
	return buff.Bytes()
}

// BuildHeader 生成编码后的Header用于发送
func (que *StarQueue) BuildHeader(length uint32) []byte {
	var buff bytes.Buffer
	que.Msgid++
	buff.Write(header)
	buff.Write(Uint32ToByte(length))
	buff.Write(Uint16ToByte(que.Msgid))
	return buff.Bytes()
}

type unFinMsg struct {
	ID         uint16
	LengthRecv uint32
	// HeaderMsg 信息头，应当为14位：8位识别码+4位长度码+2位id
	HeaderMsg []byte
	RecvMsg   []byte
}

func (que *StarQueue) push2list(msg MsgQueue) {
	que.MsgPool <- msg
}

// ParseMessage 用于解析收到的msg信息
func (que *StarQueue) ParseMessage(msg []byte, conn interface{}) error {
	return que.parseMessage(msg, conn)
}

// parseMessage 用于解析收到的msg信息
func (que *StarQueue) parseMessage(msg []byte, conn interface{}) error {
	tmp, ok := que.UnFinMsg.Load(conn)
	if ok { //存在未完成的信息
		lastMsg := tmp.(*unFinMsg)
		headerLen := len(lastMsg.HeaderMsg)
		if headerLen < 14 { //未完成头标题
			//传输的数据不能填充header头
			if len(msg) < 14-headerLen {
				//加入header头并退出
				lastMsg.HeaderMsg = bytesMerge(lastMsg.HeaderMsg, msg)
				que.UnFinMsg.Store(conn, lastMsg)
				return nil
			}
			//获取14字节完整的header
			header := msg[0 : 14-headerLen]
			lastMsg.HeaderMsg = bytesMerge(lastMsg.HeaderMsg, header)
			//检查收到的header是否为认证header
			//若不是，丢弃并重新来过
			if !checkHeader(lastMsg.HeaderMsg[0:8]) {
				que.UnFinMsg.Delete(conn)
				if len(msg) == 0 {
					return nil
				}
				return que.parseMessage(msg, conn)
			}
			//获得本数据包长度
			lastMsg.LengthRecv = ByteToUint32(lastMsg.HeaderMsg[8:12])
			//获得本数据包ID
			lastMsg.ID = ByteToUint16(lastMsg.HeaderMsg[12:14])
			//存入列表
			que.UnFinMsg.Store(conn, lastMsg)
			msg = msg[14-headerLen:]
			if uint32(len(msg)) < lastMsg.LengthRecv {
				lastMsg.RecvMsg = msg
				que.UnFinMsg.Store(conn, lastMsg)
				return nil
			}
			if uint32(len(msg)) >= lastMsg.LengthRecv {
				lastMsg.RecvMsg = msg[0:lastMsg.LengthRecv]
				if que.Encode {
					lastMsg.RecvMsg = que.DecodeFunc(lastMsg.RecvMsg)
				}
				msg = msg[lastMsg.LengthRecv:]
				storeMsg := MsgQueue{
					ID:   lastMsg.ID,
					Msg:  lastMsg.RecvMsg,
					Conn: conn,
				}
				//que.restoreMu.Lock()
				que.push2list(storeMsg)
				//que.restoreMu.Unlock()
				que.UnFinMsg.Delete(conn)
				return que.parseMessage(msg, conn)
			}
		} else {
			lastID := int(lastMsg.LengthRecv) - len(lastMsg.RecvMsg)
			if lastID < 0 {
				que.UnFinMsg.Delete(conn)
				return que.parseMessage(msg, conn)
			}
			if len(msg) >= lastID {
				lastMsg.RecvMsg = bytesMerge(lastMsg.RecvMsg, msg[0:lastID])
				if que.Encode {
					lastMsg.RecvMsg = que.DecodeFunc(lastMsg.RecvMsg)
				}
				storeMsg := MsgQueue{
					ID:   lastMsg.ID,
					Msg:  lastMsg.RecvMsg,
					Conn: conn,
				}
				que.push2list(storeMsg)
				que.UnFinMsg.Delete(conn)
				if len(msg) == lastID {
					return nil
				}
				msg = msg[lastID:]
				return que.parseMessage(msg, conn)
			}
			lastMsg.RecvMsg = bytesMerge(lastMsg.RecvMsg, msg)
			que.UnFinMsg.Store(conn, lastMsg)
			return nil
		}
	}
	if len(msg) == 0 {
		return nil
	}
	var start int
	if start = searchHeader(msg); start == -1 {
		return errors.New("data format error")
	}
	msg = msg[start:]
	lastMsg := unFinMsg{}
	que.UnFinMsg.Store(conn, &lastMsg)
	return que.parseMessage(msg, conn)
}

func checkHeader(msg []byte) bool {
	if len(msg) != 8 {
		return false
	}
	for k, v := range msg {
		if v != header[k] {
			return false
		}
	}
	return true
}

func searchHeader(msg []byte) int {
	if len(msg) < 8 {
		return 0
	}
	for k, v := range msg {
		find := 0
		if v == header[0] {
			for k2, v2 := range header {
				if msg[k+k2] == v2 {
					find++
				} else {
					break
				}
			}
			if find == 8 {
				return k
			}
		}
	}
	return -1
}

func bytesMerge(src ...[]byte) []byte {
	var buff bytes.Buffer
	for _, v := range src {
		buff.Write(v)
	}
	return buff.Bytes()
}

// Restore 获取收到的信息
func (que *StarQueue) Restore() (MsgQueue, error) {
	if que.duration.Seconds() == 0 {
		que.duration = 86400 * time.Second
	}
	for {
		select {
		case <-que.ctx.Done():
			return MsgQueue{}, errors.New("Stoped By External Function Call")
		case <-time.After(que.duration):
			if que.duration != 0 {
				return MsgQueue{}, os.ErrDeadlineExceeded
			}
		case data, ok := <-que.MsgPool:
			if !ok {
				return MsgQueue{}, os.ErrClosed
			}
			return data, nil
		}
	}
}

// RestoreOne 获取收到的一个信息
//兼容性修改
func (que *StarQueue) RestoreOne() (MsgQueue, error) {
	return que.Restore()
}

// Stop 立即停止Restore
func (que *StarQueue) Stop() {
	que.cancel()
}

// RestoreDuration Restore最大超时时间
func (que *StarQueue) RestoreDuration(tm time.Duration) {
	que.duration = tm
}

func (que *StarQueue) RestoreChan() <-chan MsgQueue {
	return que.MsgPool
}