github.com/m3db/m3@v1.5.1-0.20231129193456-75a402aa583b/src/x/pool/checked_bytes_test.go (about)

     1  // Copyright (c) 2016 Uber Technologies, Inc.
     2  //
     3  // Permission is hereby granted, free of charge, to any person obtaining a copy
     4  // of this software and associated documentation files (the "Software"), to deal
     5  // in the Software without restriction, including without limitation the rights
     6  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
     7  // copies of the Software, and to permit persons to whom the Software is
     8  // furnished to do so, subject to the following conditions:
     9  //
    10  // The above copyright notice and this permission notice shall be included in
    11  // all copies or substantial portions of the Software.
    12  //
    13  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    14  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    15  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    16  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    17  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    18  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    19  // THE SOFTWARE.
    20  
    21  package pool
    22  
    23  import (
    24  	"testing"
    25  
    26  	"github.com/stretchr/testify/assert"
    27  )
    28  
    29  func TestCheckedBytesPool(t *testing.T) {
    30  	p := getCheckedBytesPool(2, []int{5, 10})
    31  	p.Init()
    32  
    33  	assert.Equal(t, 2, checkedBytesPoolBucketLen(p, 0))
    34  	assert.Equal(t, 2, checkedBytesPoolBucketLen(p, 1))
    35  
    36  	b1 := p.Get(1)
    37  	b1.IncRef()
    38  
    39  	assert.Equal(t, 1, checkedBytesPoolBucketLen(p, 0))
    40  	assert.Equal(t, 2, checkedBytesPoolBucketLen(p, 1))
    41  
    42  	assert.Equal(t, 0, b1.Len())
    43  	assert.Equal(t, 5, b1.Cap())
    44  	b1.Append('a')
    45  
    46  	b2 := p.Get(3)
    47  	b2.IncRef()
    48  
    49  	assert.Equal(t, 0, checkedBytesPoolBucketLen(p, 0))
    50  	assert.Equal(t, 2, checkedBytesPoolBucketLen(p, 1))
    51  
    52  	assert.Equal(t, 0, b2.Len())
    53  	assert.Equal(t, 5, b2.Cap())
    54  	b2.Append('b')
    55  
    56  	assert.NotEqual(t, b1.Bytes(), b2.Bytes())
    57  
    58  	copiedB1 := append([]byte(nil), b1.Bytes()...)
    59  	b1.DecRef()
    60  	b1.Finalize()
    61  
    62  	assert.Equal(t, 1, checkedBytesPoolBucketLen(p, 0))
    63  	assert.Equal(t, 2, checkedBytesPoolBucketLen(p, 1))
    64  
    65  	b3 := p.Get(2)
    66  	b3.IncRef()
    67  	assert.Equal(t, 0, b3.Len())
    68  	assert.Equal(t, 5, b3.Cap())
    69  	assert.Equal(t, copiedB1, b3.Bytes()[:1])
    70  }
    71  
    72  func TestAppendByteChecked(t *testing.T) {
    73  	p := getCheckedBytesPool(1, []int{3, 10})
    74  	p.Init()
    75  	vals := []byte{'a', 'b', 'c', 'd', 'e', 'f', 'g', 'h', 'i'}
    76  
    77  	b1 := p.Get(2)
    78  	b1.IncRef()
    79  
    80  	firstB1 := b1
    81  
    82  	for _, val := range vals {
    83  		if b, swapped := AppendByteChecked(b1, val, p); swapped {
    84  			b1.DecRef()
    85  			b1.Finalize()
    86  			b1 = b
    87  			b1.IncRef()
    88  		}
    89  	}
    90  
    91  	// Ensure swapped out with new pooled bytes
    92  	assert.NotEqual(t, firstB1, b1)
    93  
    94  	assert.Equal(t, vals, b1.Bytes())
    95  	assert.Equal(t, 9, b1.Len())
    96  	assert.Equal(t, 10, b1.Cap())
    97  
    98  	// Ensure reusing the first slice we retrieved
    99  	b2 := p.Get(2)
   100  	b2.IncRef()
   101  	assert.Equal(t, 3, b2.Cap())
   102  	assert.Equal(t, b1.Bytes()[:3], b2.Bytes()[:3])
   103  }
   104  
   105  func getCheckedBytesPool(
   106  	bucketSizes int,
   107  	bucketCaps []int,
   108  ) *checkedBytesPool {
   109  	buckets := make([]Bucket, len(bucketCaps))
   110  	for i, cap := range bucketCaps {
   111  		buckets[i] = Bucket{
   112  			Count:    Size(bucketSizes),
   113  			Capacity: cap,
   114  		}
   115  	}
   116  
   117  	return NewCheckedBytesPool(buckets, nil, func(s []Bucket) BytesPool {
   118  		return NewBytesPool(s, nil)
   119  	}).(*checkedBytesPool)
   120  }
   121  
   122  func checkedBytesPoolBucketLen(
   123  	p *checkedBytesPool,
   124  	bucket int,
   125  ) int {
   126  	bucketizedPool := p.pool.(*bucketizedObjectPool)
   127  	objectPool := bucketizedPool.buckets[bucket].pool.(*objectPool)
   128  	return len(objectPool.values)
   129  }