Skip to content

Commit c0ebd07

Browse files
authored
Migrate dapr/utils/byteslicepool to kit (dapr#73)
No code changes (aside from those needed to appease the linter) Signed-off-by: ItalyPaleAle <[email protected]>
1 parent 2d30434 commit c0ebd07

File tree

2 files changed

+144
-0
lines changed

2 files changed

+144
-0
lines changed

byteslicepool/byteslicepool.go

Lines changed: 85 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,85 @@
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+
}
Lines changed: 59 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,59 @@
1+
/*
2+
Copyright 2023 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+
"testing"
18+
19+
"github.com/stretchr/testify/assert"
20+
)
21+
22+
func TestByteSlicePool(t *testing.T) {
23+
minCap := 32
24+
pool := NewByteSlicePool(minCap)
25+
26+
bs := pool.Get(minCap)
27+
assert.Empty(t, bs)
28+
assert.Equal(t, minCap, cap(bs))
29+
30+
pool.Put(bs)
31+
bs2 := pool.Get(minCap)
32+
assert.Equal(t, &bs, &bs2)
33+
assert.Equal(t, minCap, cap(bs2))
34+
35+
for i := 0; i < minCap; i++ {
36+
bs2 = append(bs2, 0)
37+
}
38+
39+
// Less than minCap
40+
// Capacity will not change after resize
41+
size2 := 16
42+
bs2 = pool.Resize(bs2, size2)
43+
assert.Equal(t, size2, len(bs2)) //nolint:testifylint
44+
assert.Equal(t, minCap, cap(bs2))
45+
46+
// Less than twice the minCap
47+
// Will automatically expand to twice the original capacity
48+
size3 := 48
49+
bs2 = pool.Resize(bs2, size3)
50+
assert.Equal(t, size3, len(bs2)) //nolint:testifylint
51+
assert.Equal(t, minCap*2, cap(bs2))
52+
53+
// More than twice the minCap
54+
// Will automatically expand to the specified size
55+
size4 := 128
56+
bs2 = pool.Resize(bs2, size4)
57+
assert.Equal(t, size4, len(bs2)) //nolint:testifylint
58+
assert.Equal(t, size4, cap(bs2))
59+
}

0 commit comments

Comments
 (0)