ftr: add vitess utils

container/gxbucketpool/bucketpool.go

@@ -0,0 +1,124 @@
+/*
+Copyright 2019 The Vitess 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.
+*/
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 gxbucketpool
+
+import (
+	"math"
+	"sync"
+)
+
+type sizedPool struct {
+	size int
+	pool sync.Pool
+}
+
+func newSizedPool(size int) *sizedPool {
+	return &sizedPool{
+		size: size,
+		pool: sync.Pool{
+			New: func() interface{} { return makeSlicePointer(size) },
+		},
+	}
+}
+
+// Pool is actually multiple pools which store buffers of specific size.
+// i.e. it can be three pools which return buffers 32K, 64K and 128K.
+type Pool struct {
+	minSize int
+	maxSize int
+	pools   []*sizedPool
+}
+
+// New returns Pool which has buckets from minSize to maxSize.
+// Buckets increase with the power of two, i.e with multiplier 2: [2b, 4b, 16b, ... , 1024b]
+// Last pool will always be capped to maxSize.
+func New(minSize, maxSize int) *Pool {
+	if maxSize < minSize {
+		panic("maxSize can't be less than minSize")
+	}
+	const multiplier = 2
+	var pools []*sizedPool
+	curSize := minSize
+	for curSize < maxSize {
+		pools = append(pools, newSizedPool(curSize))
+		curSize *= multiplier
+	}
+	pools = append(pools, newSizedPool(maxSize))
+	return &Pool{
+		minSize: minSize,
+		maxSize: maxSize,
+		pools:   pools,
+	}
+}
+
+func (p *Pool) findPool(size int) *sizedPool {
+	if size > p.maxSize {
+		return nil
+	}
+	idx := int(math.Ceil(math.Log2(float64(size) / float64(p.minSize))))
+	if idx < 0 {
+		idx = 0
+	}
+	if idx > len(p.pools)-1 {
+		return nil
+	}
+	return p.pools[idx]
+}
+
+// Get returns pointer to []byte which has len size.
+// If there is no bucket with buffers >= size, slice will be allocated.
+func (p *Pool) Get(size int) *[]byte {
+	sp := p.findPool(size)
+	if sp == nil {
+		return makeSlicePointer(size)
+	}
+	buf := sp.pool.Get().(*[]byte)
+	*buf = (*buf)[:size]
+	return buf
+}
+
+// Put returns pointer to slice to some bucket. Discards slice for which there is no bucket
+func (p *Pool) Put(b *[]byte) {
+	sp := p.findPool(cap(*b))
+	if sp == nil {
+		return
+	}
+	*b = (*b)[:cap(*b)]
+	sp.pool.Put(b)
+}
+
+func makeSlicePointer(size int) *[]byte {
+	data := make([]byte, size)
+	return &data
+}

container/gxbucketpool/bucketpool_test.go

@@ -0,0 +1,235 @@
+/*
+Copyright 2019 The Vitess 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.
+*/
+
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You 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 gxbucketpool
+
+import (
+	"math/rand"
+	"testing"
+)
+
+func TestPool(t *testing.T) {
+	maxSize := 16384
+	pool := New(1024, maxSize)
+	if pool.maxSize != maxSize {
+		t.Fatalf("Invalid max pool size: %d, expected %d", pool.maxSize, maxSize)
+	}
+	if len(pool.pools) != 5 {
+		t.Fatalf("Invalid number of pools: %d, expected %d", len(pool.pools), 5)
+	}
+
+	buf := pool.Get(64)
+	if len(*buf) != 64 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 1024 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+
+	// get from same pool, check that length is right
+	buf = pool.Get(128)
+	if len(*buf) != 128 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 1024 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	// get boundary size
+	buf = pool.Get(1024)
+	if len(*buf) != 1024 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 1024 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	// get from the middle
+	buf = pool.Get(5000)
+	if len(*buf) != 5000 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 8192 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	// check last pool
+	buf = pool.Get(16383)
+	if len(*buf) != 16383 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 16384 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	// get big buffer
+	buf = pool.Get(16385)
+	if len(*buf) != 16385 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 16385 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+}
+
+func TestPoolOneSize(t *testing.T) {
+	maxSize := 1024
+	pool := New(1024, maxSize)
+	if pool.maxSize != maxSize {
+		t.Fatalf("Invalid max pool size: %d, expected %d", pool.maxSize, maxSize)
+	}
+	buf := pool.Get(64)
+	if len(*buf) != 64 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 1024 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	buf = pool.Get(1025)
+	if len(*buf) != 1025 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 1025 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+}
+
+func TestPoolTwoSizeNotMultiplier(t *testing.T) {
+	maxSize := 2000
+	pool := New(1024, maxSize)
+	if pool.maxSize != maxSize {
+		t.Fatalf("Invalid max pool size: %d, expected %d", pool.maxSize, maxSize)
+	}
+	buf := pool.Get(64)
+	if len(*buf) != 64 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 1024 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	buf = pool.Get(2001)
+	if len(*buf) != 2001 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 2001 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+}
+
+func TestPoolWeirdMaxSize(t *testing.T) {
+	maxSize := 15000
+	pool := New(1024, maxSize)
+	if pool.maxSize != maxSize {
+		t.Fatalf("Invalid max pool size: %d, expected %d", pool.maxSize, maxSize)
+	}
+
+	buf := pool.Get(14000)
+	if len(*buf) != 14000 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 15000 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+
+	buf = pool.Get(16383)
+	if len(*buf) != 16383 {
+		t.Fatalf("unexpected buf length: %d", len(*buf))
+	}
+	if cap(*buf) != 16383 {
+		t.Fatalf("unexpected buf cap: %d", cap(*buf))
+	}
+	pool.Put(buf)
+}
+
+func TestFuzz(t *testing.T) {
+	maxTestSize := 16384
+	for i := 0; i < 20000; i++ {
+		minSize := rand.Intn(maxTestSize)
+		maxSize := rand.Intn(maxTestSize-minSize) + minSize
+		p := New(minSize, maxSize)
+		bufSize := rand.Intn(maxTestSize)
+		buf := p.Get(bufSize)
+		if len(*buf) != bufSize {
+			t.Fatalf("Invalid length %d, expected %d", len(*buf), bufSize)
+		}
+		sPool := p.findPool(bufSize)
+		if sPool == nil {
+			if cap(*buf) != len(*buf) {
+				t.Fatalf("Invalid cap %d, expected %d", cap(*buf), len(*buf))
+			}
+		} else {
+			if cap(*buf) != sPool.size {
+				t.Fatalf("Invalid cap %d, expected %d", cap(*buf), sPool.size)
+			}
+		}
+		p.Put(buf)
+	}
+}
+
+func BenchmarkPool(b *testing.B) {
+	pool := New(2, 16384)
+	b.SetParallelism(16)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			randomSize := rand.Intn(pool.maxSize)
+			data := pool.Get(randomSize)
+			pool.Put(data)
+		}
+	})
+}
+
+func BenchmarkPoolGet(b *testing.B) {
+	pool := New(2, 16384)
+	b.SetParallelism(16)
+	b.ResetTimer()
+	b.RunParallel(func(pb *testing.PB) {
+		for pb.Next() {
+			randomSize := rand.Intn(pool.maxSize)
+			data := pool.Get(randomSize)
+			_ = data
+		}
+	})
+}