github.com/clubpay/ronykit/kit@v0.14.4-0.20240515065620-d0dace45cbc7/utils/batch/batcher_test.go (about)

     1  package batch_test
     2  
     3  import (
     4  	"fmt"
     5  	"math/rand"
     6  	"sync"
     7  	"sync/atomic"
     8  	"testing"
     9  	"time"
    10  
    11  	"github.com/clubpay/ronykit/kit/utils/batch"
    12  	. "github.com/onsi/ginkgo/v2"
    13  	. "github.com/onsi/gomega"
    14  )
    15  
    16  func TestGateway(t *testing.T) {
    17  	RegisterFailHandler(Fail)
    18  
    19  	RunSpecs(t, "RonyKit/Kit/Utils Suite")
    20  }
    21  
    22  var _ = Describe("Flusher Without WaitTime", func() {
    23  	var out, in int64
    24  	f := batch.NewMulti[int, batch.NA](
    25  		func(targetID string, entries []batch.Entry[int, batch.NA]) {
    26  			time.Sleep(time.Millisecond * 100)
    27  			atomic.AddInt64(&out, int64(len(entries)))
    28  		},
    29  		batch.WithBatchSize(20),
    30  		batch.WithMaxWorkers(10),
    31  	)
    32  
    33  	wg := sync.WaitGroup{}
    34  	total := int64(10000)
    35  	for i := 0; i < int(total); i++ {
    36  		wg.Add(1)
    37  		go func() {
    38  			f.EnterAndWait(
    39  				fmt.Sprintf("T%d", rand.Intn(3)),
    40  				batch.NewEntry[int, batch.NA](rand.Intn(10), nil),
    41  			)
    42  			atomic.AddInt64(&in, 1)
    43  			wg.Done()
    44  		}()
    45  	}
    46  	wg.Wait()
    47  	It("should flush all entries", func() {
    48  		for _, q := range f.Pool() {
    49  			Expect(q.EntryChan()).To(BeEmpty())
    50  		}
    51  		Expect(in).To(Equal(total))
    52  		Expect(out).To(Equal(total))
    53  	})
    54  })
    55  
    56  var _ = Describe("Flusher With WaitTime", func() {
    57  	var out, in int64
    58  	f := batch.NewMulti[int, batch.NA](
    59  		func(targetID string, entries []batch.Entry[int, batch.NA]) {
    60  			time.Sleep(time.Millisecond * 100)
    61  			atomic.AddInt64(&out, int64(len(entries)))
    62  			for _, e := range entries {
    63  				e.Value()
    64  			}
    65  		},
    66  		batch.WithBatchSize(20),
    67  		batch.WithMaxWorkers(10),
    68  		batch.WithMinWaitTime(250*time.Millisecond),
    69  	)
    70  
    71  	wg := sync.WaitGroup{}
    72  	total := int64(10000)
    73  	for i := 0; i < int(total); i++ {
    74  		wg.Add(1)
    75  		go func() {
    76  			f.EnterAndWait(
    77  				fmt.Sprintf("T%d", rand.Intn(3)),
    78  				batch.NewEntry[int, batch.NA](rand.Intn(10), nil),
    79  			)
    80  			atomic.AddInt64(&in, 1)
    81  			wg.Done()
    82  		}()
    83  	}
    84  	wg.Wait()
    85  
    86  	It("should flush all entries", func() {
    87  		for _, q := range f.Pool() {
    88  			Expect(q.EntryChan()).To(BeEmpty())
    89  		}
    90  		Expect(in).To(Equal(total))
    91  		Expect(out).To(Equal(total))
    92  	})
    93  })
    94  
    95  var _ = Describe("Flusher With Callback", func() {
    96  	var out, in int64
    97  	f := batch.NewMulti[int, int](
    98  		func(targetID string, entries []batch.Entry[int, int]) {
    99  			time.Sleep(time.Millisecond * 100)
   100  			atomic.AddInt64(&out, int64(len(entries)))
   101  			for _, e := range entries {
   102  				e.Callback(e.Value())
   103  			}
   104  		},
   105  		batch.WithBatchSize(20),
   106  		batch.WithMaxWorkers(10),
   107  		batch.WithMinWaitTime(250*time.Millisecond),
   108  	)
   109  
   110  	wg := sync.WaitGroup{}
   111  	total := int64(10000)
   112  	var sum int64
   113  	for i := 0; i < int(total); i++ {
   114  		wg.Add(1)
   115  		go func(x int) {
   116  			f.EnterAndWait(
   117  				"sameID",
   118  				batch.NewEntry(
   119  					x,
   120  					func(out int) { atomic.AddInt64(&sum, int64(out)) },
   121  				),
   122  			)
   123  			atomic.AddInt64(&in, 1)
   124  			wg.Done()
   125  		}(i)
   126  	}
   127  	wg.Wait()
   128  
   129  	It("should flush all entries", func() {
   130  		for _, q := range f.Pool() {
   131  			Expect(q.EntryChan()).To(BeEmpty())
   132  		}
   133  		Expect(in).To(Equal(total))
   134  		Expect(out).To(Equal(total))
   135  		Expect(sum).To(Equal(total * (total - 1) / 2))
   136  	})
   137  })