mirror of https://github.com/dapr/kit.git
Migrate dapr/utils/byteslicepool to kit (#73)
No code changes (aside from those needed to appease the linter) Signed-off-by: ItalyPaleAle <43508+ItalyPaleAle@users.noreply.github.com>
This commit is contained in:
parent
2d30434d91
commit
c0ebd07f3a
|
@ -0,0 +1,85 @@
|
|||
/*
|
||||
Copyright 2021 The Dapr 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 byteslicepool
|
||||
|
||||
import (
|
||||
"sync"
|
||||
)
|
||||
|
||||
/*
|
||||
Originally based on https://github.com/xdg-go/zzz-slice-recycling
|
||||
Copyright (C) 2019 by David A. Golden
|
||||
License (Apache2): https://github.com/xdg-go/zzz-slice-recycling/blob/master/LICENSE
|
||||
*/
|
||||
|
||||
// ByteSlicePool is a wrapper around sync.Pool to get []byte objects with a given capacity.
|
||||
type ByteSlicePool struct {
|
||||
MinCap int
|
||||
pool *sync.Pool
|
||||
}
|
||||
|
||||
// NewByteSlicePool returns a new ByteSlicePool object.
|
||||
func NewByteSlicePool(minCap int) *ByteSlicePool {
|
||||
return &ByteSlicePool{
|
||||
MinCap: minCap,
|
||||
pool: &sync.Pool{},
|
||||
}
|
||||
}
|
||||
|
||||
// Get a slice from the pool.
|
||||
// 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.
|
||||
func (sp ByteSlicePool) Get(capacity int) []byte {
|
||||
bp := sp.pool.Get()
|
||||
if bp == nil {
|
||||
if capacity < sp.MinCap {
|
||||
capacity = sp.MinCap
|
||||
}
|
||||
return make([]byte, 0, capacity)
|
||||
}
|
||||
buf := bp.([]byte)
|
||||
// This will be optimized by the compiler
|
||||
for i := range buf {
|
||||
buf[i] = 0
|
||||
}
|
||||
return buf[:0]
|
||||
}
|
||||
|
||||
// Put a slice back in the pool.
|
||||
func (sp ByteSlicePool) Put(bs []byte) {
|
||||
// The linter here complains because we're putting a slice rather than a pointer in the pool.
|
||||
// The complain is valid, because doing so does cause an allocation for the local copy of the slice header.
|
||||
// However, this is ok for us because given how we use ByteSlicePool, we can't keep around the pointer we took out.
|
||||
// See this thread for some discussion: https://github.com/dominikh/go-tools/issues/1336
|
||||
//nolint:staticcheck
|
||||
sp.pool.Put(bs)
|
||||
}
|
||||
|
||||
// Resize a byte slice, making sure that it has enough capacity for a given size.
|
||||
func (sp ByteSlicePool) Resize(orig []byte, size int) []byte {
|
||||
if size < cap(orig) {
|
||||
return orig[0:size]
|
||||
}
|
||||
|
||||
// Allocate a new byte slice and then discard the old one, too small, so it can be garbage collected
|
||||
temp := make([]byte, size, max(size, cap(orig)*2))
|
||||
copy(temp, orig)
|
||||
return temp
|
||||
}
|
||||
|
||||
func max(x, y int) int {
|
||||
if x < y {
|
||||
return y
|
||||
}
|
||||
return x
|
||||
}
|
|
@ -0,0 +1,59 @@
|
|||
/*
|
||||
Copyright 2023 The Dapr 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 byteslicepool
|
||||
|
||||
import (
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
)
|
||||
|
||||
func TestByteSlicePool(t *testing.T) {
|
||||
minCap := 32
|
||||
pool := NewByteSlicePool(minCap)
|
||||
|
||||
bs := pool.Get(minCap)
|
||||
assert.Empty(t, bs)
|
||||
assert.Equal(t, minCap, cap(bs))
|
||||
|
||||
pool.Put(bs)
|
||||
bs2 := pool.Get(minCap)
|
||||
assert.Equal(t, &bs, &bs2)
|
||||
assert.Equal(t, minCap, cap(bs2))
|
||||
|
||||
for i := 0; i < minCap; i++ {
|
||||
bs2 = append(bs2, 0)
|
||||
}
|
||||
|
||||
// Less than minCap
|
||||
// Capacity will not change after resize
|
||||
size2 := 16
|
||||
bs2 = pool.Resize(bs2, size2)
|
||||
assert.Equal(t, size2, len(bs2)) //nolint:testifylint
|
||||
assert.Equal(t, minCap, cap(bs2))
|
||||
|
||||
// Less than twice the minCap
|
||||
// Will automatically expand to twice the original capacity
|
||||
size3 := 48
|
||||
bs2 = pool.Resize(bs2, size3)
|
||||
assert.Equal(t, size3, len(bs2)) //nolint:testifylint
|
||||
assert.Equal(t, minCap*2, cap(bs2))
|
||||
|
||||
// More than twice the minCap
|
||||
// Will automatically expand to the specified size
|
||||
size4 := 128
|
||||
bs2 = pool.Resize(bs2, size4)
|
||||
assert.Equal(t, size4, len(bs2)) //nolint:testifylint
|
||||
assert.Equal(t, size4, cap(bs2))
|
||||
}
|
Loading…
Reference in New Issue