github.com/cloudwego/kitex@v0.9.0/pkg/remote/trans/nphttp2/buffer.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 nphttp2
    18  
    19  import (
    20  	"net"
    21  	"sync"
    22  
    23  	"github.com/bytedance/gopkg/lang/mcache"
    24  
    25  	"github.com/cloudwego/kitex/pkg/remote"
    26  )
    27  
    28  const (
    29  	// min is 4k
    30  	minMallocSize = 4 * 1024
    31  )
    32  
    33  // GRPCConn implement WriteFrame and ReadFrame
    34  type GRPCConn interface {
    35  	net.Conn
    36  	// WriteFrame set header and data buffer into frame with nocopy
    37  	WriteFrame(hdr, data []byte) (n int, err error)
    38  	// ReadFrame read a frame and return header and payload data
    39  	ReadFrame() (hdr, data []byte, err error)
    40  }
    41  
    42  type buffer struct {
    43  	conn       GRPCConn
    44  	rbuf       []byte // for read
    45  	whdr, wbuf []byte // for write
    46  }
    47  
    48  var (
    49  	_ remote.ByteBuffer = (*buffer)(nil)
    50  	_ remote.FrameWrite = (*buffer)(nil)
    51  )
    52  
    53  var bufferPool = sync.Pool{
    54  	New: func() interface{} {
    55  		return &buffer{}
    56  	},
    57  }
    58  
    59  func newBuffer(conn GRPCConn) *buffer {
    60  	buf := bufferPool.Get().(*buffer)
    61  	buf.conn = conn
    62  	return buf
    63  }
    64  
    65  func (b *buffer) growRbuf(n int) {
    66  	capacity := cap(b.rbuf)
    67  	if capacity >= n {
    68  		return
    69  	}
    70  	if n < minMallocSize {
    71  		n = minMallocSize
    72  	}
    73  	buf := mcache.Malloc(0, n)
    74  	if cap(b.rbuf) > 0 {
    75  		mcache.Free(b.rbuf)
    76  	}
    77  	b.rbuf = buf
    78  }
    79  
    80  func (b *buffer) Next(n int) (p []byte, err error) {
    81  	b.growRbuf(n)
    82  	_, err = b.conn.Read(b.rbuf[:n])
    83  	if err != nil {
    84  		return nil, err
    85  	}
    86  	return b.rbuf[:n], nil
    87  }
    88  
    89  func (b *buffer) WriteHeader(buf []byte) (err error) {
    90  	b.whdr = buf
    91  	return nil
    92  }
    93  
    94  func (b *buffer) WriteData(buf []byte) (err error) {
    95  	b.wbuf = buf
    96  	return nil
    97  }
    98  
    99  func (b *buffer) Flush() (err error) {
   100  	_, err = b.conn.WriteFrame(b.whdr, b.wbuf)
   101  	b.whdr = nil
   102  	b.wbuf = nil
   103  	return err
   104  }
   105  
   106  func (b *buffer) Release(e error) (err error) {
   107  	if cap(b.rbuf) > 0 {
   108  		mcache.Free(b.rbuf)
   109  	}
   110  	b.conn = nil
   111  	b.rbuf = nil
   112  	b.whdr = nil
   113  	b.wbuf = nil
   114  	bufferPool.Put(b)
   115  	return e
   116  }
   117  
   118  // === unimplemented ===
   119  
   120  func (b *buffer) Read(p []byte) (n int, err error) {
   121  	panic("implement me")
   122  }
   123  
   124  func (b *buffer) Write(p []byte) (n int, err error) {
   125  	panic("implement me")
   126  }
   127  
   128  func (b *buffer) Peek(n int) (buf []byte, err error) {
   129  	panic("implement me")
   130  }
   131  
   132  func (b *buffer) Skip(n int) (err error) {
   133  	panic("implement me")
   134  }
   135  
   136  func (b *buffer) ReadableLen() (n int) {
   137  	panic("implement me")
   138  }
   139  
   140  func (b *buffer) ReadLen() (n int) {
   141  	panic("implement me")
   142  }
   143  
   144  func (b *buffer) ReadString(n int) (s string, err error) {
   145  	panic("implement me")
   146  }
   147  
   148  func (b *buffer) ReadBinary(n int) (p []byte, err error) {
   149  	panic("implement me")
   150  }
   151  
   152  func (b *buffer) Malloc(n int) (buf []byte, err error) {
   153  	panic("implement me")
   154  }
   155  
   156  func (b *buffer) MallocLen() (length int) {
   157  	panic("implement me")
   158  }
   159  
   160  func (b *buffer) WriteString(s string) (n int, err error) {
   161  	panic("implement me")
   162  }
   163  
   164  func (b *buffer) WriteBinary(data []byte) (n int, err error) {
   165  	panic("implement me")
   166  }
   167  
   168  func (b *buffer) NewBuffer() remote.ByteBuffer {
   169  	panic("implement me")
   170  }
   171  
   172  func (b *buffer) AppendBuffer(buf remote.ByteBuffer) (err error) {
   173  	panic("implement me")
   174  }
   175  
   176  func (b *buffer) Bytes() (buf []byte, err error) {
   177  	panic("implement me")
   178  }