nps/lib/enet/conn.go

package enet

import (
	"ehang.io/nps/lib/pool"
	"errors"
	"net"
	"sync"
	"syscall"
)

type Conn interface {
	net.Conn
	Reset(int) error
	Clear()
	Readable() bool
	AllBytes() ([]byte, error)
	SyscallConn() (syscall.RawConn, error)
}

var _ Conn = (*ReaderConn)(nil)

var bp = pool.NewBufferPool(MaxReadSize)

const MaxReadSize = 32 * 1024

// ReaderConn is an implement of reusable data connection
type ReaderConn struct {
	buf      []byte
	nowIndex int
	hasRead  int
	hasClear bool
	net.Conn
	sync.RWMutex
}

// NewReaderConn returns a new ReaderConn
func NewReaderConn(conn net.Conn) *ReaderConn {
	return &ReaderConn{Conn: conn, buf: bp.Get()}
}

// SyscallConn returns a raw network connection
func (rc *ReaderConn) SyscallConn() (syscall.RawConn, error) {
	return rc.Conn.(syscall.Conn).SyscallConn()
}

// Read is an implement of Net.Conn Read function
func (rc *ReaderConn) Read(b []byte) (n int, err error) {
	rc.Lock()
	defer rc.Unlock()
	if rc.hasClear || (rc.nowIndex == rc.hasRead && rc.hasRead == MaxReadSize) {
		if !rc.hasClear {
			rc.Clear()
		}
		return rc.Conn.Read(b)
	}
	if rc.hasRead > rc.nowIndex {
		n = copy(b, rc.buf[rc.nowIndex:rc.hasRead])
		rc.nowIndex += n
		return
	}
	if rc.hasRead == MaxReadSize {
		n = copy(b, rc.buf[rc.nowIndex:rc.hasRead])
		rc.nowIndex += n
		return
	}
	err = rc.readOnce()
	if err != nil {
		return
	}
	n = copy(b, rc.buf[rc.nowIndex:rc.hasRead])
	rc.nowIndex += n
	return
}

// readOnce
func (rc *ReaderConn) readOnce() error {
	// int(math.Min(float64(MaxReadSize-rc.hasRead), float64(len(b)-(rc.hasRead-rc.nowIndex))))
	// read as much as possible to judge whether there is still readable
	n, err := rc.Conn.Read(rc.buf[rc.nowIndex : rc.hasRead+MaxReadSize-rc.hasRead])
	rc.hasRead += n
	return err
}

// Readable return whether there is data in the buffer
func (rc *ReaderConn) Readable() bool {
	return (rc.hasRead - rc.nowIndex) > 0
}

// AllBytes return all data in the buffer
func (rc *ReaderConn) AllBytes() ([]byte, error) {
	rc.Lock()
	defer rc.Unlock()
	if rc.hasRead == 0 {
		if err := rc.readOnce(); err != nil {
			return nil, err
		}
	}
	if !rc.Readable() {
		return nil, errors.New("can not read '")
	}
	b := rc.buf[rc.nowIndex:rc.hasRead]
	rc.nowIndex = rc.hasRead
	return b, nil
}

// Reset will reset data index
func (rc *ReaderConn) Reset(n int) error {
	if !rc.hasClear {
		rc.nowIndex = n
		return nil
	}
	return errors.New("the enet can not reset anymore")
}

// Clear will put buf to pool and can not reuse anymore
func (rc *ReaderConn) Clear() {
	rc.hasClear = true
	bp.Put(rc.buf)
}