github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/netpoll/bytebuf.go

/*
 * Copyright 2021 CloudWeGo Authors
 *
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

package netpoll

import (
	"errors"
	"io"
	"sync"

	"github.com/cloudwego/netpoll"

	"github.com/cloudwego/kitex/pkg/remote"
)

var bytebufPool sync.Pool

func init() {
	bytebufPool.New = newNetpollByteBuffer
}

// NewReaderByteBuffer creates a new remote.ByteBuffer using the given netpoll.ZeroCopyReader.
func NewReaderByteBuffer(r netpoll.Reader) remote.ByteBuffer {
	bytebuf := bytebufPool.Get().(*netpollByteBuffer)
	bytebuf.reader = r
	// TODO(wangtieju): fix me when netpoll support netpoll.Reader
	// and LinkBuffer not support io.Reader, type assertion would fail when r is from NewBuffer
	if ir, ok := r.(io.Reader); ok {
		bytebuf.ioReader = ir
	}
	bytebuf.status = remote.BitReadable
	bytebuf.readSize = 0
	return bytebuf
}

// NewWriterByteBuffer creates a new remote.ByteBuffer using the given netpoll.ZeroCopyWriter.
func NewWriterByteBuffer(w netpoll.Writer) remote.ByteBuffer {
	bytebuf := bytebufPool.Get().(*netpollByteBuffer)
	bytebuf.writer = w
	// TODO(wangtieju): fix me when netpoll support netpoll.Writer
	// and LinkBuffer not support io.Reader, type assertion would fail when w is from NewBuffer
	if iw, ok := w.(io.Writer); ok {
		bytebuf.ioWriter = iw
	}
	bytebuf.status = remote.BitWritable
	return bytebuf
}

// NewReaderWriterByteBuffer creates a new remote.ByteBuffer using the given netpoll.ZeroCopyReadWriter.
func NewReaderWriterByteBuffer(rw netpoll.ReadWriter) remote.ByteBuffer {
	bytebuf := bytebufPool.Get().(*netpollByteBuffer)
	bytebuf.writer = rw
	bytebuf.reader = rw
	// TODO(wangtieju): fix me when netpoll support netpoll.ReadWriter
	// and LinkBuffer not support io.ReadWriter, type assertion would fail when rw is from NewBuffer
	if irw, ok := rw.(io.ReadWriter); ok {
		bytebuf.ioReader = irw
		bytebuf.ioWriter = irw
	}
	bytebuf.status = remote.BitWritable | remote.BitReadable
	return bytebuf
}

func newNetpollByteBuffer() interface{} {
	return &netpollByteBuffer{}
}

type netpollByteBuffer struct {
	writer   netpoll.Writer
	reader   netpoll.Reader
	ioReader io.Reader
	ioWriter io.Writer
	status   int
	readSize int
}

var _ remote.ByteBuffer = &netpollByteBuffer{}

// Next reads n bytes sequentially, returns the original address.
func (b *netpollByteBuffer) Next(n int) (p []byte, err error) {
	if b.status&remote.BitReadable == 0 {
		return nil, errors.New("unreadable buffer, cannot support Next")
	}
	if p, err = b.reader.Next(n); err == nil {
		b.readSize += n
	}
	return
}

// Peek returns the next n bytes without advancing the reader.
func (b *netpollByteBuffer) Peek(n int) (buf []byte, err error) {
	if b.status&remote.BitReadable == 0 {
		return nil, errors.New("unreadable buffer, cannot support Peek")
	}
	return b.reader.Peek(n)
}

// Skip is used to skip n bytes, it's much faster than Next.
// Skip will not cause release.
func (b *netpollByteBuffer) Skip(n int) (err error) {
	if b.status&remote.BitReadable == 0 {
		return errors.New("unreadable buffer, cannot support Skip")
	}
	return b.reader.Skip(n)
}

// ReadableLen returns the total length of readable buffer.
func (b *netpollByteBuffer) ReadableLen() (n int) {
	if b.status&remote.BitReadable == 0 {
		return -1
	}
	return b.reader.Len()
}

// Read implement io.Reader
func (b *netpollByteBuffer) Read(p []byte) (n int, err error) {
	if b.status&remote.BitReadable == 0 {
		return -1, errors.New("unreadable buffer, cannot support Read")
	}
	if b.ioReader != nil {
		return b.ioReader.Read(p)
	}
	return -1, errors.New("ioReader is nil")
}

// ReadString is a more efficient way to read string than Next.
func (b *netpollByteBuffer) ReadString(n int) (s string, err error) {
	if b.status&remote.BitReadable == 0 {
		return "", errors.New("unreadable buffer, cannot support ReadString")
	}
	if s, err = b.reader.ReadString(n); err == nil {
		b.readSize += n
	}
	return
}

// ReadBinary like ReadString.
// Returns a copy of original buffer.
func (b *netpollByteBuffer) ReadBinary(n int) (p []byte, err error) {
	if b.status&remote.BitReadable == 0 {
		return p, errors.New("unreadable buffer, cannot support ReadBinary")
	}
	if p, err = b.reader.ReadBinary(n); err == nil {
		b.readSize += n
	}
	return
}

// Malloc n bytes sequentially in the writer buffer.
func (b *netpollByteBuffer) Malloc(n int) (buf []byte, err error) {
	if b.status&remote.BitWritable == 0 {
		return nil, errors.New("unwritable buffer, cannot support Malloc")
	}
	return b.writer.Malloc(n)
}

// MallocLen returns the total length of the buffer malloced.
func (b *netpollByteBuffer) MallocLen() (length int) {
	if b.status&remote.BitWritable == 0 {
		return -1
	}
	return b.writer.MallocLen()
}

// Write implement io.Writer
func (b *netpollByteBuffer) Write(p []byte) (n int, err error) {
	if b.status&remote.BitWritable == 0 {
		return -1, errors.New("unwritable buffer, cannot support Write")
	}
	if b.ioWriter != nil {
		return b.ioWriter.Write(p)
	}
	return -1, errors.New("ioWriter is nil")
}

// WriteString is a more efficient way to write string, using the unsafe method to convert the string to []byte.
func (b *netpollByteBuffer) WriteString(s string) (n int, err error) {
	if b.status&remote.BitWritable == 0 {
		return -1, errors.New("unwritable buffer, cannot support WriteString")
	}
	return b.writer.WriteString(s)
}

// WriteBinary writes the []byte directly. Callers must guarantee that the []byte doesn't change.
func (b *netpollByteBuffer) WriteBinary(p []byte) (n int, err error) {
	if b.status&remote.BitWritable == 0 {
		return -1, errors.New("unwritable buffer, cannot support WriteBinary")
	}
	return b.writer.WriteBinary(p)
}

// WriteDirect is a way to write []byte without copying, and splits the original buffer.
func (b *netpollByteBuffer) WriteDirect(p []byte, remainCap int) error {
	if b.status&remote.BitWritable == 0 {
		return errors.New("unwritable buffer, cannot support WriteBinary")
	}
	return b.writer.WriteDirect(p, remainCap)
}

// ReadLen returns the size already read.
func (b *netpollByteBuffer) ReadLen() (n int) {
	return b.readSize
}

// Flush writes any malloc data to the underlying io.Writer.
// The malloced buffer must be set correctly.
func (b *netpollByteBuffer) Flush() (err error) {
	if b.status&remote.BitWritable == 0 {
		return errors.New("unwritable buffer, cannot support Flush")
	}
	return b.writer.Flush()
}

// NewBuffer returns a new writable remote.ByteBuffer.
func (b *netpollByteBuffer) NewBuffer() remote.ByteBuffer {
	return NewWriterByteBuffer(netpoll.NewLinkBuffer())
}

// AppendBuffer appends buf to the original buffer.
func (b *netpollByteBuffer) AppendBuffer(buf remote.ByteBuffer) (err error) {
	subBuf := buf.(*netpollByteBuffer)
	err = b.writer.Append(subBuf.writer)
	buf.Release(nil)
	return
}

// Bytes are not supported in netpoll bytebuf.
func (b *netpollByteBuffer) Bytes() (buf []byte, err error) {
	return nil, errors.New("method Bytes() not support in netpoll bytebuf")
}

// Release will free the buffer already read.
// After release, buffer read by Next/Skip/Peek is invalid.
func (b *netpollByteBuffer) Release(e error) (err error) {
	if b.reader != nil {
		// 重复执行Release nil panic
		err = b.reader.Release()
	}
	b.zero()
	bytebufPool.Put(b)
	return
}

func (b *netpollByteBuffer) zero() {
	b.writer = nil
	b.reader = nil
	b.ioReader = nil
	b.ioWriter = nil
	b.status = 0
	b.readSize = 0
}