github.com/thanos-io/thanos@v0.32.5/pkg/pool/pool_test.go (about)

     1  // Copyright (c) The Thanos Authors.
     2  // Licensed under the Apache License 2.0.
     3  
     4  package pool
     5  
     6  import (
     7  	"strings"
     8  	"sync"
     9  	"testing"
    10  	"time"
    11  
    12  	"github.com/pkg/errors"
    13  	"go.uber.org/goleak"
    14  
    15  	"github.com/efficientgo/core/testutil"
    16  )
    17  
    18  func TestMain(m *testing.M) {
    19  	goleak.VerifyTestMain(m)
    20  }
    21  
    22  func TestBytesPool(t *testing.T) {
    23  	chunkPool, err := NewBucketedBytes(10, 100, 2, 1000)
    24  	testutil.Ok(t, err)
    25  
    26  	testutil.Equals(t, []int{10, 20, 40, 80}, chunkPool.sizes)
    27  
    28  	for i := 0; i < 10; i++ {
    29  		b, err := chunkPool.Get(40)
    30  		testutil.Ok(t, err)
    31  
    32  		testutil.Equals(t, uint64(40), chunkPool.usedTotal)
    33  
    34  		if i%2 == 0 {
    35  			for j := 0; j < 6; j++ {
    36  				*b = append(*b, []byte{'1', '2', '3', '4', '5'}...)
    37  			}
    38  		}
    39  		chunkPool.Put(b)
    40  	}
    41  
    42  	for i := 0; i < 10; i++ {
    43  		b, err := chunkPool.Get(19)
    44  		testutil.Ok(t, err)
    45  		chunkPool.Put(b)
    46  	}
    47  
    48  	// Outside of any bucket.
    49  	b, err := chunkPool.Get(1000)
    50  	testutil.Ok(t, err)
    51  	chunkPool.Put(b)
    52  
    53  	// Check size limitation.
    54  	b1, err := chunkPool.Get(500)
    55  	testutil.Ok(t, err)
    56  
    57  	b2, err := chunkPool.Get(600)
    58  	testutil.NotOk(t, err)
    59  	testutil.Equals(t, ErrPoolExhausted, err)
    60  
    61  	chunkPool.Put(b1)
    62  	chunkPool.Put(b2)
    63  
    64  	testutil.Equals(t, uint64(0), chunkPool.usedTotal)
    65  }
    66  
    67  func TestRacePutGet(t *testing.T) {
    68  	chunkPool, err := NewBucketedBytes(3, 100, 2, 5000)
    69  	testutil.Ok(t, err)
    70  
    71  	s := sync.WaitGroup{}
    72  
    73  	const goroutines = 100
    74  
    75  	// Start multiple goroutines: they always Get and Put two byte slices
    76  	// to which they write their contents and check if the data is still
    77  	// there after writing it, before putting it back.
    78  	errs := make(chan error, goroutines)
    79  	stop := make(chan struct{})
    80  
    81  	f := func(txt string, grow bool) {
    82  		defer s.Done()
    83  		for {
    84  			select {
    85  			case <-stop:
    86  				return
    87  			default:
    88  				c, err := chunkPool.Get(len(txt))
    89  				if err != nil {
    90  					errs <- errors.Wrapf(err, "goroutine %s", txt)
    91  					return
    92  				}
    93  
    94  				*c = append(*c, txt...)
    95  				if string(*c) != txt {
    96  					errs <- errors.New("expected to get the data just written")
    97  					return
    98  				}
    99  				if grow {
   100  					*c = append(*c, txt...)
   101  					*c = append(*c, txt...)
   102  					if string(*c) != txt+txt+txt {
   103  						errs <- errors.New("expected to get the data just written")
   104  						return
   105  					}
   106  				}
   107  
   108  				chunkPool.Put(c)
   109  			}
   110  		}
   111  	}
   112  
   113  	for i := 0; i < goroutines; i++ {
   114  		s.Add(1)
   115  		// make sure we start multiple goroutines with same len buf requirements, to hit same pools
   116  		s := strings.Repeat(string(byte(i)), i%10)
   117  		// some of the goroutines will append more elements to the provided slice
   118  		grow := i%2 == 0
   119  		go f(s, grow)
   120  	}
   121  
   122  	time.Sleep(1 * time.Second)
   123  	close(stop)
   124  	s.Wait()
   125  	select {
   126  	case err := <-errs:
   127  		testutil.Ok(t, err)
   128  	default:
   129  	}
   130  }