google.golang.org/grpc@v1.62.1/shared_buffer_pool.go (about)

     1  /*
     2   *
     3   * Copyright 2023 gRPC authors.
     4   *
     5   * Licensed under the Apache License, Version 2.0 (the "License");
     6   * you may not use this file except in compliance with the License.
     7   * You may obtain a copy of the License at
     8   *
     9   *     http://www.apache.org/licenses/LICENSE-2.0
    10   *
    11   * Unless required by applicable law or agreed to in writing, software
    12   * distributed under the License is distributed on an "AS IS" BASIS,
    13   * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    14   * See the License for the specific language governing permissions and
    15   * limitations under the License.
    16   *
    17   */
    18  
    19  package grpc
    20  
    21  import "sync"
    22  
    23  // SharedBufferPool is a pool of buffers that can be shared, resulting in
    24  // decreased memory allocation. Currently, in gRPC-go, it is only utilized
    25  // for parsing incoming messages.
    26  //
    27  // # Experimental
    28  //
    29  // Notice: This API is EXPERIMENTAL and may be changed or removed in a
    30  // later release.
    31  type SharedBufferPool interface {
    32  	// Get returns a buffer with specified length from the pool.
    33  	//
    34  	// The returned byte slice may be not zero initialized.
    35  	Get(length int) []byte
    36  
    37  	// Put returns a buffer to the pool.
    38  	Put(*[]byte)
    39  }
    40  
    41  // NewSharedBufferPool creates a simple SharedBufferPool with buckets
    42  // of different sizes to optimize memory usage. This prevents the pool from
    43  // wasting large amounts of memory, even when handling messages of varying sizes.
    44  //
    45  // # Experimental
    46  //
    47  // Notice: This API is EXPERIMENTAL and may be changed or removed in a
    48  // later release.
    49  func NewSharedBufferPool() SharedBufferPool {
    50  	return &simpleSharedBufferPool{
    51  		pools: [poolArraySize]simpleSharedBufferChildPool{
    52  			newBytesPool(level0PoolMaxSize),
    53  			newBytesPool(level1PoolMaxSize),
    54  			newBytesPool(level2PoolMaxSize),
    55  			newBytesPool(level3PoolMaxSize),
    56  			newBytesPool(level4PoolMaxSize),
    57  			newBytesPool(0),
    58  		},
    59  	}
    60  }
    61  
    62  // simpleSharedBufferPool is a simple implementation of SharedBufferPool.
    63  type simpleSharedBufferPool struct {
    64  	pools [poolArraySize]simpleSharedBufferChildPool
    65  }
    66  
    67  func (p *simpleSharedBufferPool) Get(size int) []byte {
    68  	return p.pools[p.poolIdx(size)].Get(size)
    69  }
    70  
    71  func (p *simpleSharedBufferPool) Put(bs *[]byte) {
    72  	p.pools[p.poolIdx(cap(*bs))].Put(bs)
    73  }
    74  
    75  func (p *simpleSharedBufferPool) poolIdx(size int) int {
    76  	switch {
    77  	case size <= level0PoolMaxSize:
    78  		return level0PoolIdx
    79  	case size <= level1PoolMaxSize:
    80  		return level1PoolIdx
    81  	case size <= level2PoolMaxSize:
    82  		return level2PoolIdx
    83  	case size <= level3PoolMaxSize:
    84  		return level3PoolIdx
    85  	case size <= level4PoolMaxSize:
    86  		return level4PoolIdx
    87  	default:
    88  		return levelMaxPoolIdx
    89  	}
    90  }
    91  
    92  const (
    93  	level0PoolMaxSize = 16                     //  16  B
    94  	level1PoolMaxSize = level0PoolMaxSize * 16 // 256  B
    95  	level2PoolMaxSize = level1PoolMaxSize * 16 //   4 KB
    96  	level3PoolMaxSize = level2PoolMaxSize * 16 //  64 KB
    97  	level4PoolMaxSize = level3PoolMaxSize * 16 //   1 MB
    98  )
    99  
   100  const (
   101  	level0PoolIdx = iota
   102  	level1PoolIdx
   103  	level2PoolIdx
   104  	level3PoolIdx
   105  	level4PoolIdx
   106  	levelMaxPoolIdx
   107  	poolArraySize
   108  )
   109  
   110  type simpleSharedBufferChildPool interface {
   111  	Get(size int) []byte
   112  	Put(any)
   113  }
   114  
   115  type bufferPool struct {
   116  	sync.Pool
   117  
   118  	defaultSize int
   119  }
   120  
   121  func (p *bufferPool) Get(size int) []byte {
   122  	bs := p.Pool.Get().(*[]byte)
   123  
   124  	if cap(*bs) < size {
   125  		p.Pool.Put(bs)
   126  
   127  		return make([]byte, size)
   128  	}
   129  
   130  	return (*bs)[:size]
   131  }
   132  
   133  func newBytesPool(size int) simpleSharedBufferChildPool {
   134  	return &bufferPool{
   135  		Pool: sync.Pool{
   136  			New: func() any {
   137  				bs := make([]byte, size)
   138  				return &bs
   139  			},
   140  		},
   141  		defaultSize: size,
   142  	}
   143  }
   144  
   145  // nopBufferPool is a buffer pool just makes new buffer without pooling.
   146  type nopBufferPool struct {
   147  }
   148  
   149  func (nopBufferPool) Get(length int) []byte {
   150  	return make([]byte, length)
   151  }
   152  
   153  func (nopBufferPool) Put(*[]byte) {
   154  }