From c0ebd07f3aef15867b7fac0f73d6e67d4ef85c3c Mon Sep 17 00:00:00 2001
From: "Alessandro (Ale) Segala" <43508+ItalyPaleAle@users.noreply.github.com>
Date: Mon, 30 Oct 2023 11:55:45 -0700
Subject: [PATCH] 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>
---
 byteslicepool/byteslicepool.go      | 85 +++++++++++++++++++++++++++++
 byteslicepool/byteslicepool_test.go | 59 ++++++++++++++++++++
 2 files changed, 144 insertions(+)
 create mode 100644 byteslicepool/byteslicepool.go
 create mode 100644 byteslicepool/byteslicepool_test.go

diff --git a/byteslicepool/byteslicepool.go b/byteslicepool/byteslicepool.go
new file mode 100644
index 0000000..ba7b69d
--- /dev/null
+++ b/byteslicepool/byteslicepool.go
@@ -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
+}
diff --git a/byteslicepool/byteslicepool_test.go b/byteslicepool/byteslicepool_test.go
new file mode 100644
index 0000000..c8f4e71
--- /dev/null
+++ b/byteslicepool/byteslicepool_test.go
@@ -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))
+}