api: add Allocator interface for response buffers

Previously, response buffer allocation was hardcoded with an internal
pooler implementation. Users had no control over memory management
for responses, which could be problematic for applications with
specific memory requirements or custom allocation strategies.

The new Allocator interface allows users to implement custom buffer
allocation and reuse strategies:

```
// Allocator is an interface for allocating and deallocating byte
// slices.
type Allocator interface {
    // Get returns a pointer to a byte slice of at least the given
    // length.
    // The caller should not assume anything about the slice's
    // capacity.
    //
    // If the allocator cannot allocate a buffer (e.g., invalid
    // length), it returns nil. The caller must handle this case
    // appropriately.
    Get(length int) *[]byte
    // Put returns a byte slice to the allocator for reuse.
    // After calling Put, the caller must not use the slice.
    //
    // The caller must ensure that the slice length remains unchanged
    // between Get and Put calls. Modifying the slice length before
    // calling Put may prevent the allocator from properly reusing the
    // buffer.
    Put(buf *[]byte)
}
```

The PoolAllocator type provides a ready-to-use implementation based
on sync.Pool for power-of-two sized byte slices.

The Opts.Allocator option enables configuring a custom allocator
for a connection. This is useful for applications that need to
optimize memory usage or integrate with custom memory management
systems.

Closes #493

Author: oleg-jukovec

CHANGELOG.md

@@ -23,6 +23,11 @@ Versioning](http://semver.org/spec/v2.0.0.html) except to the first release.
 * Resources allocated for a `Future` object created by the `Connection` type
   could be released with the `Future.Release()` call.
 * Added function String() for type interval (#322).
+* New `Allocator` interface for custom allocation of response buffers (#493).
+* New `PoolAllocator` type that implements `Allocator` using sync.Pool for
+  power-of-two sized byte slices (#493).
+* New `Opts.Allocator` option to configure a custom allocator for a
+  connection (#493).
 
 ### Changed
 

alloc.go

@@ -0,0 +1,18 @@
+package tarantool
+
+// Allocator is an interface for allocating and deallocating byte slices.
+type Allocator interface {
+	// Get returns a pointer to a byte slice of at least the given length.
+	// The caller should not assume anything about the slice's capacity.
+	//
+	// If the allocator cannot allocate a buffer (e.g., invalid length), it
+	// returns nil. The caller must handle this case appropriately.
+	Get(length int) *[]byte
+	// Put returns a byte slice to the allocator for reuse.
+	// After calling Put, the caller must not use the slice.
+	//
+	// The caller must ensure that the slice length remains unchanged between
+	// Get and Put calls. Modifying the slice length before calling Put may
+	// prevent the allocator from properly reusing the buffer.
+	Put(buf *[]byte)
+}

connection.go

@@ -211,7 +211,8 @@ type Connection struct {
 	// requestCnt is a counter of active requests.
 	requestCnt int64
 
-	pool allocator
+	// alloc is an allocator for response buffers.
+	alloc Allocator
 }
 
 var _ = Connector(&Connection{}) // Check compatibility with connector interface.
@@ -329,9 +330,9 @@ type Opts struct {
 	Handle interface{}
 	// Logger is user specified logger used for error messages.
 	Logger Logger
-	// PoolSizes is a list of pool sizes (power-of-two) for each connection,
-	// nil value means default -- not using pool inside.
-	PoolSizes []int
+	// Allocator is used to allocate and deallocate response buffers.
+	// If nil, default Go allocation is used.
+	Allocator Allocator
 }
 
 // Connect creates and configures a new Connection.
@@ -378,7 +379,12 @@ func Connect(ctx context.Context, dialer Dialer, opts Opts) (conn *Connection, e
 		conn.opts.Logger = defaultLogger{}
 	}
 
+	if conn.opts.Allocator != nil {
+		conn.alloc = conn.opts.Allocator
+	}
+
 	conn.cond = sync.NewCond(&conn.mutex)
+
 	if conn.opts.Reconnect > 0 {
 		// We don't need these mutex.Lock()/mutex.Unlock() here, but
 		// runReconnects() expects mutex.Lock() to be set, so it's
@@ -400,10 +406,6 @@ func Connect(ctx context.Context, dialer Dialer, opts Opts) (conn *Connection, e
 		go conn.timeouts()
 	}
 
-	if conn.opts.PoolSizes != nil {
-		conn.pool = newPooler(conn.opts.PoolSizes)
-	}
-
 	// TODO: reload schema after reconnect.
 	if !conn.opts.SkipSchema {
 		schema, err := GetSchema(conn)
@@ -891,7 +893,7 @@ func (conn *Connection) reader(r io.Reader, c Conn) {
 	for atomic.LoadUint32(&conn.state) != connClosed {
 		var err error
 
-		buf, err = read(r, conn.lenbuf[:], conn.pool)
+		buf, err = read(r, conn.lenbuf[:], conn.alloc)
 
 		if err != nil {
 			err = ClientError{
@@ -1225,36 +1227,40 @@ func (conn *Connection) timeouts() {
 	}
 }
 
-// read uses args to allocate slices for responses using sync.Pool.
-// data must be released later using Release.
-func read(r io.Reader, lenbuf []byte, pool allocator) (buf smallBuf, err error) {
+// read reads a response from the reader and allocates a buffer using the
+// allocator. If alloc is nil, a regular Go allocation is used.
+// The returned buffer could be released using Release.
+func read(r io.Reader, lenbuf []byte, alloc Allocator) (smallBuf, error) {
 	var length uint64
 
-	if _, err = io.ReadFull(r, lenbuf); err != nil {
-		return
+	if _, err := io.ReadFull(r, lenbuf); err != nil {
+		return smallBuf{}, fmt.Errorf("failed to read response length: %w", err)
 	}
+
 	if lenbuf[0] != 0xce {
-		err = errors.New("wrong response header")
-		return
+		return smallBuf{}, errors.New("wrong response header")
 	}
+
 	length = (uint64(lenbuf[1]) << 24) +
 		(uint64(lenbuf[2]) << 16) +
 		(uint64(lenbuf[3]) << 8) +
 		uint64(lenbuf[4])
 
 	switch {
 	case length == 0:
-		err = errors.New("response should not be 0 length")
-		return
+		return smallBuf{}, errors.New("response should not be 0 length")
 	case length > math.MaxUint32:
-		err = errors.New("response is too big")
-		return
+		return smallBuf{}, errors.New("response is too big")
 	}
 
-	buf = CreateBuf(pool, int(length))
-	_, err = io.ReadFull(r, buf.b)
+	buf := createBuf(alloc, int(length))
+	if _, err := io.ReadFull(r, buf.b); err != nil {
+		buf.Release()
 
-	return
+		return smallBuf{}, fmt.Errorf("failed to read response: %w", err)
+	}
+
+	return buf, nil
 }
 
 func (conn *Connection) nextRequestId(context bool) (requestId uint32) {

example_custom_request_test.go

@@ -221,7 +221,7 @@ func ExampleRequest_Response_manual() {
 	defer func() { _ = conn.Close() }()
 
 	// Insert test data.
-	_, err := conn.Do(tarantool.NewInsertRequest(spaceNo).
+	_, err := conn.Do(tarantool.NewReplaceRequest(spaceNo).
 		Tuple([]interface{}{uint(1111), "hello", "world"}),
 	).Get()
 	if err != nil {
@@ -258,7 +258,7 @@ func ExampleRequest_Response_manualDecodeTyped() {
 	defer func() { _ = conn.Close() }()
 
 	// Insert test data.
-	_, err := conn.Do(tarantool.NewInsertRequest(spaceNo).
+	_, err := conn.Do(tarantool.NewReplaceRequest(spaceNo).
 		Tuple([]interface{}{uint(1111), "hello", "world"}),
 	).Get()
 	if err != nil {

example_test.go

@@ -1542,3 +1542,50 @@ func ExampleFdDialer() {
 	// Output:
 	// <nil>
 }
+
+// ExamplePoolAllocator demonstrates how to use PoolAllocator for
+// buffer pooling to reduce memory allocations when working with
+// a connection.
+func ExamplePoolAllocator() {
+	// Create a pool allocator with pools for 256, 1024, and 4096 byte slices.
+	// Exponents must be sorted in ascending order and in range [0, 31].
+	allocator, err := tarantool.NewPoolAllocator([]int{8, 10, 12})
+	if err != nil {
+		fmt.Printf("Failed to create allocator: %s\n", err)
+		return
+	}
+
+	// Use the allocator with a connection.
+	optsWithAlloc := opts
+	optsWithAlloc.Allocator = allocator
+
+	conn := exampleConnect(dialer, optsWithAlloc)
+	defer func() { _ = conn.Close() }()
+
+	// Insert a tuple.
+	fut := conn.Do(tarantool.NewReplaceRequest("test").
+		Tuple([]interface{}{uint(1111), "hello", "world"}))
+	_, err = fut.Get()
+	if err != nil {
+		fmt.Printf("Failed to replace: %s\n", err)
+		return
+	}
+	// Release the buffer back to the pool after use.
+	fut.Release()
+
+	// Select the tuple - the allocator will be used for response buffers.
+	fut = conn.Do(tarantool.NewSelectRequest("test").
+		Index("primary").
+		Key(tarantool.UintKey{1111}))
+	data, err := fut.Get()
+	if err != nil {
+		fmt.Printf("Failed to select: %s\n", err)
+		return
+	}
+	fmt.Println(data)
+	// Release the buffer back to the pool after use.
+	fut.Release()
+
+	// Output:
+	// [[1111 hello world]]
+}

poolalloc.go

@@ -0,0 +1,119 @@
+package tarantool
+
+import (
+	"errors"
+	"math/bits"
+	"slices"
+	"sync"
+)
+
+// PoolAllocator implements the Allocator interface using a set of sync.Pool
+// instances for different power-of-two sized byte slices.
+//
+// The exponents parameter in NewPoolAllocator specifies the powers of two for
+// the pool sizes. For example, []int{8, 10, 12} creates pools for 256, 1024,
+// and 4096 byte slices.
+type PoolAllocator struct {
+	pool []*sync.Pool
+	size []int
+	help []int
+}
+
+var _ Allocator = (*PoolAllocator)(nil)
+
+// NewPoolAllocator creates a new PoolAllocator with the given exponents.
+// Each exponent represents a power of two pool size. For example, exponent 10
+// creates a pool for 1024-byte slices.
+//
+// Exponents must be sorted in ascending order and each exponent must be
+// in range [0, 31].
+func NewPoolAllocator(exponents []int) (*PoolAllocator, error) {
+	hSize := 32
+
+	for i, s := range exponents {
+		if s < 0 || s >= hSize {
+			return nil, errors.New("exponent must be in range [0, 31]")
+		}
+
+		if i > 0 && exponents[i-1] >= s {
+			return nil, errors.New("exponents must be sorted in ascending order")
+		}
+	}
+
+	var p = PoolAllocator{
+		size: make([]int, len(exponents)),
+		pool: make([]*sync.Pool, len(exponents)),
+		help: slices.Repeat([]int{-1}, hSize),
+	}
+
+	for i, s := range exponents {
+		p.size[i] = 1 << s
+		p.help[s] = i
+		p.pool[i] = &sync.Pool{
+			New: func() interface{} {
+				buf := make([]byte, p.size[i])
+				return &buf
+			},
+		}
+	}
+
+	for i := hSize - 2; i >= 0; i-- {
+		if p.help[i] != -1 {
+			continue
+		}
+
+		p.help[i] = p.help[i+1]
+	}
+
+	return &p, nil
+}
+
+func (p *PoolAllocator) getInd(size int) int {
+	if size <= 0 {
+		return -1
+	}
+
+	idx := bits.Len(uint(size - 1))
+	if idx >= len(p.help) {
+		return -1
+	}
+
+	return p.help[idx]
+}
+
+// Get returns a pointer to a byte slice of at least the given length.
+// If the requested size fits within one of the pool sizes, it returns a
+// slice from the appropriate pool.
+//
+// It returns nil if:
+//   - length is less than or equal to zero
+//   - length exceeds the maximum pool size
+func (p *PoolAllocator) Get(length int) *[]byte {
+	if length <= 0 {
+		return nil
+	}
+
+	ind := p.getInd(length)
+	if ind == -1 {
+		return nil
+	}
+
+	bs := p.pool[ind].Get().(*[]byte)
+	*bs = (*bs)[:length]
+	return bs
+}
+
+// Put returns a byte slice to the appropriate pool for reuse.
+// The slice is cleared before being returned to the pool.
+// If the slice size doesn't match any pool, it is discarded.
+// If buf is nil, the method does nothing.
+func (p *PoolAllocator) Put(buf *[]byte) {
+	if buf == nil {
+		return
+	}
+
+	if ind := p.getInd(len(*buf)); ind != -1 {
+		clear((*buf)[:len(*buf)])
+		p.pool[ind].Put(buf)
+	}
+}