github.com/cloudwego/kitex@v0.9.0/pkg/remote/bytebuf.go (about)

     1  /*
     2   * Copyright 2021 CloudWeGo Authors
     3   *
     4   * Licensed under the Apache License, Version 2.0 (the "License");
     5   * you may not use this file except in compliance with the License.
     6   * You may obtain a copy of the License at
     7   *
     8   *     http://www.apache.org/licenses/LICENSE-2.0
     9   *
    10   * Unless required by applicable law or agreed to in writing, software
    11   * distributed under the License is distributed on an "AS IS" BASIS,
    12   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13   * See the License for the specific language governing permissions and
    14   * limitations under the License.
    15   */
    16  
    17  package remote
    18  
    19  import (
    20  	"io"
    21  	"net"
    22  )
    23  
    24  // ByteBufferFactory is used to create ByteBuffer.
    25  type ByteBufferFactory interface {
    26  	NewByteBuffer(conn net.Conn) (ByteBuffer, error)
    27  }
    28  
    29  // NocopyWrite is to write []byte without copying, and splits the original buffer.
    30  // It is used with linked buffer implement.
    31  type NocopyWrite interface {
    32  	// the buf will be wrapped as a new data node no copy, then insert into the linked buffer.
    33  	// remainCap is the remain capacity of origin buff.
    34  	WriteDirect(buf []byte, remainCap int) error
    35  }
    36  
    37  // FrameWrite is to write header and data buffer separately to avoid memory copy
    38  type FrameWrite interface {
    39  	// WriteHeader set header buffer without copy
    40  	WriteHeader(buf []byte) (err error)
    41  	// WriteData set data buffer without copy
    42  	WriteData(buf []byte) (err error)
    43  }
    44  
    45  // ByteBuffer is the core abstraction of buffer in Kitex.
    46  type ByteBuffer interface {
    47  	io.ReadWriter
    48  
    49  	// Next reads the next n bytes sequentially and returns the original buffer.
    50  	Next(n int) (p []byte, err error)
    51  
    52  	// Peek returns the next n bytes without advancing the reader.
    53  	Peek(n int) (buf []byte, err error)
    54  
    55  	// Skip is used to skip the next few bytes quickly. It's faster than Next and doesn't cause release.
    56  	Skip(n int) (err error)
    57  
    58  	// Release will free the buffer. After release, buffer read by Next/Skip/Peek is invalid.
    59  	// Param e is used when the buffer release depend on error.
    60  	// For example, usually the write buffer will be released inside flush,
    61  	// but if flush error happen, write buffer may need to be released explicitly.
    62  	Release(e error) (err error)
    63  
    64  	// ReadableLen returns the total length of readable buffer.
    65  	// Return: -1 means unreadable.
    66  	ReadableLen() (n int)
    67  
    68  	// ReadLen returns the size already read.
    69  	ReadLen() (n int)
    70  
    71  	// ReadString is a more efficient way to read string than Next.
    72  	ReadString(n int) (s string, err error)
    73  
    74  	// ReadBinary like ReadString.
    75  	// Returns a copy of original buffer.
    76  	ReadBinary(n int) (p []byte, err error)
    77  
    78  	// Malloc n bytes sequentially in the writer buffer.
    79  	Malloc(n int) (buf []byte, err error)
    80  
    81  	// MallocLen returns the total length of the buffer malloced.
    82  	MallocLen() (length int)
    83  
    84  	// WriteString is a more efficient way to write string, using the unsafe method to convert the string to []byte.
    85  	WriteString(s string) (n int, err error)
    86  
    87  	// WriteBinary writes the []byte directly. Callers must guarantee that the []byte doesn't change.
    88  	WriteBinary(b []byte) (n int, err error)
    89  
    90  	// Flush writes any malloc data to the underlying io.Writer.
    91  	// The malloced buffer must be set correctly.
    92  	Flush() (err error)
    93  
    94  	// NewBuffer returns a new writable remote.ByteBuffer.
    95  	NewBuffer() ByteBuffer
    96  	// AppendBuffer appends buf to the original buffer.
    97  	AppendBuffer(buf ByteBuffer) (err error)
    98  
    99  	// Bytes return the backing bytes slice of this buffer
   100  	Bytes() (buf []byte, err error)
   101  }
   102  
   103  // ByteBufferIO wrap ByteBuffer to implement io.ReadWriter
   104  type ByteBufferIO struct {
   105  	buffer ByteBuffer
   106  }
   107  
   108  // NewByteBufferIO wraps ByBuffer to io.ReadWriter
   109  func NewByteBufferIO(buffer ByteBuffer) io.ReadWriter {
   110  	return &ByteBufferIO{
   111  		buffer: buffer,
   112  	}
   113  }
   114  
   115  // Write implements the io.ReadWriter interface.
   116  func (p *ByteBufferIO) Write(b []byte) (int, error) {
   117  	return p.buffer.WriteBinary(b)
   118  }
   119  
   120  // Read implements the io.ReadWriter interface.
   121  func (p *ByteBufferIO) Read(b []byte) (n int, err error) {
   122  	var buf []byte
   123  	readable := p.buffer.ReadableLen()
   124  	if readable == 0 {
   125  		return 0, io.EOF
   126  	} else if len(b) <= readable {
   127  		buf, err = p.buffer.Next(len(b))
   128  		if err != nil {
   129  			return
   130  		}
   131  		n = len(b)
   132  	} else {
   133  		buf, err = p.buffer.Next(readable)
   134  		if err != nil {
   135  			return
   136  		}
   137  		n = readable
   138  	}
   139  	copy(b, buf)
   140  	return
   141  }