github.com/jackc/pgx/v5@v5.5.5/internal/iobufpool/iobufpool_test.go (about)

     1  package iobufpool_test
     2  
     3  import (
     4  	"testing"
     5  
     6  	"github.com/jackc/pgx/v5/internal/iobufpool"
     7  	"github.com/stretchr/testify/assert"
     8  )
     9  
    10  func TestGetCap(t *testing.T) {
    11  	tests := []struct {
    12  		requestedLen int
    13  		expectedCap  int
    14  	}{
    15  		{requestedLen: 0, expectedCap: 256},
    16  		{requestedLen: 128, expectedCap: 256},
    17  		{requestedLen: 255, expectedCap: 256},
    18  		{requestedLen: 256, expectedCap: 256},
    19  		{requestedLen: 257, expectedCap: 512},
    20  		{requestedLen: 511, expectedCap: 512},
    21  		{requestedLen: 512, expectedCap: 512},
    22  		{requestedLen: 513, expectedCap: 1024},
    23  		{requestedLen: 1023, expectedCap: 1024},
    24  		{requestedLen: 1024, expectedCap: 1024},
    25  		{requestedLen: 33554431, expectedCap: 33554432},
    26  		{requestedLen: 33554432, expectedCap: 33554432},
    27  
    28  		// Above 32 MiB skip the pool and allocate exactly the requested size.
    29  		{requestedLen: 33554433, expectedCap: 33554433},
    30  	}
    31  	for _, tt := range tests {
    32  		buf := iobufpool.Get(tt.requestedLen)
    33  		assert.Equalf(t, tt.requestedLen, len(*buf), "bad len for requestedLen: %d", len(*buf), tt.requestedLen)
    34  		assert.Equalf(t, tt.expectedCap, cap(*buf), "bad cap for requestedLen: %d", tt.requestedLen)
    35  	}
    36  }
    37  
    38  func TestPutHandlesWrongSizedBuffers(t *testing.T) {
    39  	for putBufSize := range []int{0, 1, 128, 250, 256, 257, 1023, 1024, 1025, 1 << 28} {
    40  		putBuf := make([]byte, putBufSize)
    41  		iobufpool.Put(&putBuf)
    42  
    43  		tests := []struct {
    44  			requestedLen int
    45  			expectedCap  int
    46  		}{
    47  			{requestedLen: 0, expectedCap: 256},
    48  			{requestedLen: 128, expectedCap: 256},
    49  			{requestedLen: 255, expectedCap: 256},
    50  			{requestedLen: 256, expectedCap: 256},
    51  			{requestedLen: 257, expectedCap: 512},
    52  			{requestedLen: 511, expectedCap: 512},
    53  			{requestedLen: 512, expectedCap: 512},
    54  			{requestedLen: 513, expectedCap: 1024},
    55  			{requestedLen: 1023, expectedCap: 1024},
    56  			{requestedLen: 1024, expectedCap: 1024},
    57  			{requestedLen: 33554431, expectedCap: 33554432},
    58  			{requestedLen: 33554432, expectedCap: 33554432},
    59  
    60  			// Above 32 MiB skip the pool and allocate exactly the requested size.
    61  			{requestedLen: 33554433, expectedCap: 33554433},
    62  		}
    63  		for _, tt := range tests {
    64  			getBuf := iobufpool.Get(tt.requestedLen)
    65  			assert.Equalf(t, tt.requestedLen, len(*getBuf), "len(putBuf): %d, requestedLen: %d", len(putBuf), tt.requestedLen)
    66  			assert.Equalf(t, tt.expectedCap, cap(*getBuf), "cap(putBuf): %d, requestedLen: %d", cap(putBuf), tt.requestedLen)
    67  		}
    68  	}
    69  }
    70  
    71  func TestPutGetBufferReuse(t *testing.T) {
    72  	// There is no way to guarantee a buffer will be reused. It should be, but a GC between the Put and the Get will cause
    73  	// it not to be. So try many times.
    74  	for i := 0; i < 100000; i++ {
    75  		buf := iobufpool.Get(4)
    76  		(*buf)[0] = 1
    77  		iobufpool.Put(buf)
    78  		buf = iobufpool.Get(4)
    79  		if (*buf)[0] == 1 {
    80  			return
    81  		}
    82  	}
    83  
    84  	t.Error("buffer was never reused")
    85  }