github.com/grailbio/base@v0.0.11/gtl/tests/randomized_freepool_test.go (about)

     1  //go:generate ../generate_randomized_freepool.py --prefix=ints -DELEM=[]int --package=tests --output=int_freepool
     2  //go:generate ../generate_randomized_freepool.py --prefix=strings -DELEM=[]string --package=tests --output=string_freepool
     3  
     4  package tests
     5  
     6  import (
     7  	"sync"
     8  	"testing"
     9  
    10  	"github.com/stretchr/testify/require"
    11  )
    12  
    13  type payload = []int
    14  
    15  // Test the case where each goroutine calls Get immediately followed by Put.
    16  func TestIndependentGets(t *testing.T) {
    17  	p := NewIntsFreePool(func() payload { return []int{10, 11} }, -1)
    18  	wg := sync.WaitGroup{}
    19  	const numThreads = 100
    20  	for i := 0; i < numThreads; i++ {
    21  		wg.Add(1)
    22  		go func() {
    23  			defer wg.Done()
    24  			for i := 0; i < 10000; i++ {
    25  				v := p.Get()
    26  				require.Equal(t, []int{10, 11}, v)
    27  				p.Put(v)
    28  			}
    29  		}()
    30  	}
    31  	wg.Wait()
    32  	// Allow some slack per thread.,
    33  	require.Truef(t, p.ApproxLen() <= numThreads*2, "Pool too large: %v", p.ApproxLen())
    34  }
    35  
    36  // Test the case where each goroutine calls Get, and lets another goroutine calls Put.
    37  func TestPutsByAnotherThread(t *testing.T) {
    38  	const numThreads = 100
    39  	const getsPerThread = 1000
    40  	ch := make(chan payload, numThreads)
    41  	p := NewIntsFreePool(func() payload { return []int{20, 21} }, -1)
    42  
    43  	// Getters
    44  	getterWg := sync.WaitGroup{}
    45  	for i := 0; i < numThreads; i++ {
    46  		getterWg.Add(1)
    47  		go func() {
    48  			defer getterWg.Done()
    49  			for i := 0; i < getsPerThread; i++ {
    50  				v := p.Get()
    51  				require.Equal(t, []int{20, 21}, v)
    52  				ch <- v
    53  			}
    54  		}()
    55  	}
    56  
    57  	// Putters
    58  	putterWg := sync.WaitGroup{}
    59  	for i := 0; i < numThreads/2; i++ {
    60  		putterWg.Add(1)
    61  		go func() {
    62  			defer putterWg.Done()
    63  			for v := range ch {
    64  				require.Equal(t, []int{20, 21}, v)
    65  				p.Put(v)
    66  			}
    67  		}()
    68  	}
    69  	getterWg.Wait()
    70  	close(ch)
    71  	putterWg.Wait()
    72  	// Allow some slack
    73  	require.Truef(t, p.ApproxLen() <= numThreads*getsPerThread/20, "Pool too large: %v", p.ApproxLen())
    74  }