github.com/reusee/pr2@v0.0.0-20230630035947-72a20ff5e864/pool_test.go (about)

     1  package pr2
     2  
     3  import (
     4  	"encoding/binary"
     5  	"fmt"
     6  	"math/rand"
     7  	"runtime"
     8  	"sync"
     9  	"testing"
    10  )
    11  
    12  func TestBytesPool(t *testing.T) {
    13  	pool := NewPool(8, func() []byte {
    14  		return make([]byte, 8)
    15  	})
    16  	wg := new(sync.WaitGroup)
    17  	for i := 0; i < 200; i++ {
    18  		wg.Add(1)
    19  		i := i
    20  		go func() {
    21  			defer wg.Done()
    22  			for j := 0; j < 200; j++ {
    23  				var bs []byte
    24  				put := pool.Get(&bs)
    25  				defer put()
    26  				binary.PutUvarint(bs, uint64(i))
    27  			}
    28  		}()
    29  	}
    30  	wg.Wait()
    31  }
    32  
    33  func TestBytesPoolRC(t *testing.T) {
    34  	pool := NewPool(8, func() []byte {
    35  		return make([]byte, 8)
    36  	})
    37  	wg := new(sync.WaitGroup)
    38  	for i := 0; i < 200; i++ {
    39  		wg.Add(1)
    40  		i := i
    41  		go func() {
    42  			defer wg.Done()
    43  			for j := 0; j < 200; j++ {
    44  				var bs []byte
    45  				put, inc := pool.GetRC(&bs)
    46  				defer put()
    47  				nRef := rand.Intn(16)
    48  				for i := 0; i < nRef; i++ {
    49  					inc()
    50  				}
    51  				defer func() {
    52  					for i := 0; i < nRef; i++ {
    53  						put()
    54  					}
    55  				}()
    56  				binary.PutUvarint(bs, uint64(i))
    57  			}
    58  		}()
    59  	}
    60  	wg.Wait()
    61  }
    62  
    63  func TestBytesPoolRCOverload(t *testing.T) {
    64  	pool := NewPool(1, func() int {
    65  		return 42
    66  	})
    67  	var i int
    68  	pool.GetRC(&i)
    69  	var j int
    70  	put, inc := pool.GetRC(&j)
    71  	inc()
    72  	if put() {
    73  		t.Fatal()
    74  	}
    75  	if !put() {
    76  		t.Fatal()
    77  	}
    78  }
    79  
    80  func BenchmarkBytesPool(b *testing.B) {
    81  	pool := NewPool(8, func() []byte {
    82  		return make([]byte, 8)
    83  	})
    84  	b.ResetTimer()
    85  	for i := 0; i < b.N; i++ {
    86  		var v []byte
    87  		put := pool.Get(&v)
    88  		put()
    89  	}
    90  }
    91  
    92  func BenchmarkParallelBytesPool(b *testing.B) {
    93  	pool := NewPool(1024, func() []byte {
    94  		return make([]byte, 8)
    95  	})
    96  	b.ResetTimer()
    97  	b.RunParallel(func(pb *testing.PB) {
    98  		for pb.Next() {
    99  			var v []byte
   100  			put := pool.Get(&v)
   101  			put()
   102  		}
   103  	})
   104  }
   105  
   106  func TestGetter(t *testing.T) {
   107  	pool := NewPool(8, func() []byte {
   108  		return make([]byte, 8)
   109  	})
   110  	wg := new(sync.WaitGroup)
   111  	for i := 0; i < 200; i++ {
   112  		wg.Add(1)
   113  		i := i
   114  		go func() {
   115  			defer wg.Done()
   116  			get, put := pool.Getter()
   117  			defer put()
   118  			for j := 0; j < 200; j++ {
   119  				var v []byte
   120  				get(&v)
   121  				binary.PutUvarint(v, uint64(i))
   122  			}
   123  		}()
   124  	}
   125  	wg.Wait()
   126  }
   127  
   128  func TestPoolBadPut(t *testing.T) {
   129  	pool := NewPool(1, func() int {
   130  		return 42
   131  	})
   132  	var i int
   133  	put := pool.Get(&i)
   134  	put()
   135  	func() {
   136  		defer func() {
   137  			p := recover()
   138  			if p == nil {
   139  				t.Fatal()
   140  			}
   141  			if fmt.Sprintf("%v", p) != "bad put" {
   142  				t.Fatal()
   143  			}
   144  		}()
   145  		put()
   146  	}()
   147  }
   148  
   149  func TestPoolBadPutRC(t *testing.T) {
   150  	pool := NewPool(1, func() int {
   151  		return 42
   152  	})
   153  	var j int
   154  	pool.Get(&j)
   155  	var i int
   156  	put := pool.Get(&i)
   157  	put()
   158  	func() {
   159  		defer func() {
   160  			p := recover()
   161  			if p == nil {
   162  				t.Fatal()
   163  			}
   164  			if fmt.Sprintf("%v", p) != "bad put" {
   165  				t.Fatal()
   166  			}
   167  		}()
   168  		put()
   169  	}()
   170  }
   171  
   172  func BenchmarkPoolDrain(b *testing.B) {
   173  	pool := NewPool(uint32(runtime.NumCPU()), func() []byte {
   174  		return make([]byte, 8)
   175  	})
   176  	b.ResetTimer()
   177  	b.RunParallel(func(pb *testing.PB) {
   178  		for pb.Next() {
   179  			var v []byte
   180  			put := pool.Get(&v)
   181  			put()
   182  		}
   183  	})
   184  }
   185  
   186  func BenchmarkFastrand(b *testing.B) {
   187  	for i := 0; i < b.N; i++ {
   188  		fastrand()
   189  	}
   190  }