|
| 1 | +/* |
| 2 | +Copyright 2021 The Dapr Authors |
| 3 | +Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +you may not use this file except in compliance with the License. |
| 5 | +You may obtain a copy of the License at |
| 6 | + http://www.apache.org/licenses/LICENSE-2.0 |
| 7 | +Unless required by applicable law or agreed to in writing, software |
| 8 | +distributed under the License is distributed on an "AS IS" BASIS, |
| 9 | +WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 10 | +See the License for the specific language governing permissions and |
| 11 | +limitations under the License. |
| 12 | +*/ |
| 13 | + |
| 14 | +package byteslicepool |
| 15 | + |
| 16 | +import ( |
| 17 | + "sync" |
| 18 | +) |
| 19 | + |
| 20 | +/* |
| 21 | +Originally based on https://github.com/xdg-go/zzz-slice-recycling |
| 22 | +Copyright (C) 2019 by David A. Golden |
| 23 | +License (Apache2): https://github.com/xdg-go/zzz-slice-recycling/blob/master/LICENSE |
| 24 | +*/ |
| 25 | + |
| 26 | +// ByteSlicePool is a wrapper around sync.Pool to get []byte objects with a given capacity. |
| 27 | +type ByteSlicePool struct { |
| 28 | + MinCap int |
| 29 | + pool *sync.Pool |
| 30 | +} |
| 31 | + |
| 32 | +// NewByteSlicePool returns a new ByteSlicePool object. |
| 33 | +func NewByteSlicePool(minCap int) *ByteSlicePool { |
| 34 | + return &ByteSlicePool{ |
| 35 | + MinCap: minCap, |
| 36 | + pool: &sync.Pool{}, |
| 37 | + } |
| 38 | +} |
| 39 | + |
| 40 | +// Get a slice from the pool. |
| 41 | +// The capacity parameter is used only if we need to allocate a new byte slice; there's no guarantee a slice retrieved from the pool will have enough capacity for that. |
| 42 | +func (sp ByteSlicePool) Get(capacity int) []byte { |
| 43 | + bp := sp.pool.Get() |
| 44 | + if bp == nil { |
| 45 | + if capacity < sp.MinCap { |
| 46 | + capacity = sp.MinCap |
| 47 | + } |
| 48 | + return make([]byte, 0, capacity) |
| 49 | + } |
| 50 | + buf := bp.([]byte) |
| 51 | + // This will be optimized by the compiler |
| 52 | + for i := range buf { |
| 53 | + buf[i] = 0 |
| 54 | + } |
| 55 | + return buf[:0] |
| 56 | +} |
| 57 | + |
| 58 | +// Put a slice back in the pool. |
| 59 | +func (sp ByteSlicePool) Put(bs []byte) { |
| 60 | + // The linter here complains because we're putting a slice rather than a pointer in the pool. |
| 61 | + // The complain is valid, because doing so does cause an allocation for the local copy of the slice header. |
| 62 | + // However, this is ok for us because given how we use ByteSlicePool, we can't keep around the pointer we took out. |
| 63 | + // See this thread for some discussion: https://github.com/dominikh/go-tools/issues/1336 |
| 64 | + //nolint:staticcheck |
| 65 | + sp.pool.Put(bs) |
| 66 | +} |
| 67 | + |
| 68 | +// Resize a byte slice, making sure that it has enough capacity for a given size. |
| 69 | +func (sp ByteSlicePool) Resize(orig []byte, size int) []byte { |
| 70 | + if size < cap(orig) { |
| 71 | + return orig[0:size] |
| 72 | + } |
| 73 | + |
| 74 | + // Allocate a new byte slice and then discard the old one, too small, so it can be garbage collected |
| 75 | + temp := make([]byte, size, max(size, cap(orig)*2)) |
| 76 | + copy(temp, orig) |
| 77 | + return temp |
| 78 | +} |
| 79 | + |
| 80 | +func max(x, y int) int { |
| 81 | + if x < y { |
| 82 | + return y |
| 83 | + } |
| 84 | + return x |
| 85 | +} |
0 commit comments