github.com/cloudwego/kitex@v0.9.0/pkg/remote/default_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 remote

import (
	"errors"
	"fmt"
	"io"
	"sync"
)

// Mask bits.
const (
	BitReadable = 1 << iota
	BitWritable
)

var bytebufPool sync.Pool

func init() {
	bytebufPool.New = newDefaultByteBuffer
}

// NewWriterBuffer is used to create a defaultByteBuffer using the given size.
// NOTICE: defaultByteBuffer is only used for testing.
func NewWriterBuffer(size int) ByteBuffer {
	return newWriterByteBuffer(size)
}

// NewReaderBuffer is used to create a defaultByteBuffer using the given buf.
func NewReaderBuffer(buf []byte) ByteBuffer {
	return newReaderByteBuffer(buf)
}

// NewReaderWriterBuffer is used to create a defaultByteBuffer using the given size.
func NewReaderWriterBuffer(size int) ByteBuffer {
	return newReaderWriterByteBuffer(size)
}

type defaultByteBuffer struct {
	buff     []byte
	readIdx  int
	writeIdx int
	status   int
}

var _ ByteBuffer = &defaultByteBuffer{}

func newDefaultByteBuffer() interface{} {
	return &defaultByteBuffer{}
}

func newReaderByteBuffer(buf []byte) ByteBuffer {
	bytebuf := bytebufPool.Get().(*defaultByteBuffer)
	bytebuf.buff = buf
	bytebuf.readIdx = 0
	bytebuf.writeIdx = len(buf)
	bytebuf.status = BitReadable
	return bytebuf
}

func newWriterByteBuffer(estimatedLength int) ByteBuffer {
	if estimatedLength <= 0 {
		estimatedLength = 256 // default buffer size
	}
	bytebuf := bytebufPool.Get().(*defaultByteBuffer)
	bytebuf.buff = make([]byte, estimatedLength)
	bytebuf.status = BitWritable
	bytebuf.writeIdx = 0
	return bytebuf
}

func newReaderWriterByteBuffer(estimatedLength int) ByteBuffer {
	if estimatedLength <= 0 {
		estimatedLength = 256 // default buffer size
	}
	bytebuf := bytebufPool.Get().(*defaultByteBuffer)
	bytebuf.buff = make([]byte, estimatedLength)
	bytebuf.readIdx = 0
	bytebuf.writeIdx = 0
	bytebuf.status = BitReadable | BitWritable
	return bytebuf
}

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

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

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

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

// Read implement io.Reader
func (b *defaultByteBuffer) Read(p []byte) (n int, err error) {
	if b.status&BitReadable == 0 {
		return -1, errors.New("unreadable buffer, cannot support Read")
	}
	pl := len(p)
	var buf []byte
	readable := b.ReadableLen()
	if readable == 0 {
		return 0, io.EOF
	}
	if pl <= readable {
		buf, err = b.Next(pl)
		if err != nil {
			return
		}
		n = pl
	} else {
		buf, err = b.Next(readable)
		if err != nil {
			return
		}
		n = readable
	}
	copy(p, buf)
	return
}

// ReadString is a more efficient way to read string than Next.
func (b *defaultByteBuffer) ReadString(n int) (s string, err error) {
	if b.status&BitReadable == 0 {
		return "", errors.New("unreadable buffer, cannot support ReadString")
	}
	var buf []byte
	if buf, err = b.Next(n); err != nil {
		return "", err
	}
	return string(buf), nil
}

// ReadBinary like ReadString.
// Returns a copy of original buffer.
func (b *defaultByteBuffer) ReadBinary(n int) (p []byte, err error) {
	if b.status&BitReadable == 0 {
		return p, errors.New("unreadable buffer, cannot support ReadBinary")
	}
	var buf []byte
	if buf, err = b.Next(n); err != nil {
		return p, err
	}
	p = make([]byte, n)
	copy(p, buf)
	return p, nil
}

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

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

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

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

// WriteBinary writes the []byte into buff.
func (b *defaultByteBuffer) WriteBinary(p []byte) (n int, err error) {
	if b.status&BitWritable == 0 {
		return -1, errors.New("unwritable buffer, cannot support WriteBinary")
	}
	n = len(p)
	b.ensureWritable(n)
	copy(b.buff[b.writeIdx:b.writeIdx+n], p)
	b.writeIdx += n
	return
}

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

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

// AppendBuffer appends buf to the original buffer.
func (b *defaultByteBuffer) AppendBuffer(buf ByteBuffer) (err error) {
	subBuf := buf.(*defaultByteBuffer)
	n := subBuf.writeIdx
	b.ensureWritable(n)
	copy(b.buff[b.writeIdx:b.writeIdx+n], subBuf.buff)
	b.writeIdx += n
	buf.Release(nil)
	return
}

// Bytes is used to get the bytes written.
func (b *defaultByteBuffer) Bytes() (buf []byte, err error) {
	if b.status&BitWritable == 0 {
		return nil, errors.New("unwritable buffer, cannot support Bytes")
	}
	buf = make([]byte, b.writeIdx)
	copy(buf, b.buff[:b.writeIdx])
	return buf, nil
}

// NewBuffer returns a new writable remote.ByteBuffer.
func (b *defaultByteBuffer) NewBuffer() ByteBuffer {
	return NewWriterBuffer(256)
}

// Release will free the buffer already read.
// After release, buffer read by Next/Skip/Peek is invalid.
func (b *defaultByteBuffer) Release(e error) (err error) {
	b.zero()
	bytebufPool.Put(b)
	return
}

func (b *defaultByteBuffer) zero() {
	b.status = 0
	b.buff = nil
	b.readIdx = 0
	b.writeIdx = 0
}

func (b *defaultByteBuffer) readableCheck(n int) error {
	readableLen := b.ReadableLen()
	if b.readIdx >= b.writeIdx && n > 0 {
		return io.EOF
	}
	if readableLen < n {
		return fmt.Errorf("default byteBuffer not enough, read[%d] but remain[%d]", n, readableLen)
	}
	return nil
}

func (b *defaultByteBuffer) writableLen() (n int) {
	if b.status&BitWritable == 0 {
		return -1
	}
	return cap(b.buff) - b.writeIdx
}

func (b *defaultByteBuffer) ensureWritable(minWritableBytes int) {
	if b.writableLen() >= minWritableBytes {
		return
	}
	minNewCapacity := b.writeIdx + minWritableBytes
	newCapacity := cap(b.buff)
	for newCapacity < minNewCapacity {
		newCapacity <<= 1
	}

	buf := make([]byte, newCapacity)
	copy(buf, b.buff)
	b.buff = buf
}