github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/gonet/bytebuffer.go

/*
 * Copyright 2022 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 gonet

import (
	"errors"
	"io"
	"sync"

	"github.com/cloudwego/netpoll"

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

var rwPool sync.Pool

func init() {
	rwPool.New = newBufferReadWriter
}

var _ remote.ByteBuffer = &bufferReadWriter{}

type bufferReadWriter struct {
	reader netpoll.Reader
	writer netpoll.Writer

	ioReader io.Reader
	ioWriter io.Writer

	readSize int
	status   int
}

func newBufferReadWriter() interface{} {
	return &bufferReadWriter{}
}

// NewBufferReader creates a new remote.ByteBuffer using the given netpoll.ZeroCopyReader.
func NewBufferReader(ir io.Reader) remote.ByteBuffer {
	rw := rwPool.Get().(*bufferReadWriter)
	if npReader, ok := ir.(interface{ Reader() netpoll.Reader }); ok {
		rw.reader = npReader.Reader()
	} else {
		rw.reader = netpoll.NewReader(ir)
	}
	rw.ioReader = ir
	rw.status = remote.BitReadable
	rw.readSize = 0
	return rw
}

// NewBufferWriter creates a new remote.ByteBuffer using the given netpoll.ZeroCopyWriter.
func NewBufferWriter(iw io.Writer) remote.ByteBuffer {
	rw := rwPool.Get().(*bufferReadWriter)
	rw.writer = netpoll.NewWriter(iw)
	rw.ioWriter = iw
	rw.status = remote.BitWritable
	return rw
}

// NewBufferReadWriter creates a new remote.ByteBuffer using the given netpoll.ZeroCopyReadWriter.
func NewBufferReadWriter(irw io.ReadWriter) remote.ByteBuffer {
	rw := rwPool.Get().(*bufferReadWriter)
	rw.writer = netpoll.NewWriter(irw)
	rw.reader = netpoll.NewReader(irw)
	rw.ioWriter = irw
	rw.ioReader = irw
	rw.status = remote.BitWritable | remote.BitReadable
	return rw
}

func (rw *bufferReadWriter) readable() bool {
	return rw.status&remote.BitReadable != 0
}

func (rw *bufferReadWriter) writable() bool {
	return rw.status&remote.BitWritable != 0
}

func (rw *bufferReadWriter) Next(n int) (p []byte, err error) {
	if !rw.readable() {
		return nil, errors.New("unreadable buffer, cannot support Next")
	}
	if p, err = rw.reader.Next(n); err == nil {
		rw.readSize += n
	}
	return
}

func (rw *bufferReadWriter) Peek(n int) (buf []byte, err error) {
	if !rw.readable() {
		return nil, errors.New("unreadable buffer, cannot support Peek")
	}
	return rw.reader.Peek(n)
}

func (rw *bufferReadWriter) Skip(n int) (err error) {
	if !rw.readable() {
		return errors.New("unreadable buffer, cannot support Skip")
	}
	return rw.reader.Skip(n)
}

func (rw *bufferReadWriter) ReadableLen() (n int) {
	if !rw.readable() {
		return -1
	}
	return rw.reader.Len()
}

func (rw *bufferReadWriter) ReadString(n int) (s string, err error) {
	if !rw.readable() {
		return "", errors.New("unreadable buffer, cannot support ReadString")
	}
	if s, err = rw.reader.ReadString(n); err == nil {
		rw.readSize += n
	}
	return
}

func (rw *bufferReadWriter) ReadBinary(n int) (p []byte, err error) {
	if !rw.readable() {
		return p, errors.New("unreadable buffer, cannot support ReadBinary")
	}
	if p, err = rw.reader.ReadBinary(n); err == nil {
		rw.readSize += n
	}
	return
}

func (rw *bufferReadWriter) Read(p []byte) (n int, err error) {
	if !rw.readable() {
		return -1, errors.New("unreadable buffer, cannot support Read")
	}
	if rw.ioReader != nil {
		return rw.ioReader.Read(p)
	}
	return -1, errors.New("ioReader is nil")
}

func (rw *bufferReadWriter) ReadLen() (n int) {
	return rw.readSize
}

func (rw *bufferReadWriter) Malloc(n int) (buf []byte, err error) {
	if !rw.writable() {
		return nil, errors.New("unwritable buffer, cannot support Malloc")
	}
	return rw.writer.Malloc(n)
}

func (rw *bufferReadWriter) MallocLen() (length int) {
	if !rw.writable() {
		return -1
	}
	return rw.writer.MallocLen()
}

func (rw *bufferReadWriter) WriteString(s string) (n int, err error) {
	if !rw.writable() {
		return -1, errors.New("unwritable buffer, cannot support WriteString")
	}
	return rw.writer.WriteString(s)
}

func (rw *bufferReadWriter) WriteBinary(b []byte) (n int, err error) {
	if !rw.writable() {
		return -1, errors.New("unwritable buffer, cannot support WriteBinary")
	}
	return rw.writer.WriteBinary(b)
}

func (rw *bufferReadWriter) Flush() (err error) {
	if !rw.writable() {
		return errors.New("unwritable buffer, cannot support Flush")
	}
	return rw.writer.Flush()
}

func (rw *bufferReadWriter) Write(p []byte) (n int, err error) {
	if !rw.writable() {
		return -1, errors.New("unwritable buffer, cannot support Write")
	}
	if rw.ioWriter != nil {
		return rw.ioWriter.Write(p)
	}
	return -1, errors.New("ioWriter is nil")
}

func (rw *bufferReadWriter) Release(e error) (err error) {
	if rw.reader != nil {
		err = rw.reader.Release()
	}
	rw.zero()
	rwPool.Put(rw)
	return
}

// WriteDirect is a way to write []byte without copying, and splits the original buffer.
func (rw *bufferReadWriter) WriteDirect(p []byte, remainCap int) error {
	if !rw.writable() {
		return errors.New("unwritable buffer, cannot support WriteBinary")
	}
	return rw.writer.WriteDirect(p, remainCap)
}

func (rw *bufferReadWriter) AppendBuffer(buf remote.ByteBuffer) (err error) {
	subBuf, ok := buf.(*bufferReadWriter)
	if !ok {
		return errors.New("AppendBuffer failed, Buffer is not bufferReadWriter")
	}
	if err = rw.writer.Append(subBuf.writer); err != nil {
		return
	}
	return buf.Release(nil)
}

// NewBuffer returns a new writable remote.ByteBuffer.
func (rw *bufferReadWriter) NewBuffer() remote.ByteBuffer {
	panic("unimplemented")
}

func (rw *bufferReadWriter) Bytes() (buf []byte, err error) {
	panic("unimplemented")
}

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