LocalAI/backend/go/parakeet-cpp/batcher_test.go

package main

import (
	"sync"
	"time"

	. "github.com/onsi/ginkgo/v2"
	. "github.com/onsi/gomega"
)

var _ = Describe("batcher", func() {
	echoReply := func(reqs []*batchRequest) {
		for _, r := range reqs {
			r.reply <- batchReply{json: r.tag}
		}
	}

	It("coalesces concurrent submits into batches", func() {
		var mu sync.Mutex
		var sizes []int
		run := func(reqs []*batchRequest) {
			mu.Lock()
			sizes = append(sizes, len(reqs))
			mu.Unlock()
			echoReply(reqs)
		}
		b := newBatcher(4, 50*time.Millisecond, run)
		stop := make(chan struct{})
		go b.run(stop)
		defer close(stop)

		const N = 4
		var wg sync.WaitGroup
		got := make([]string, N)
		for i := 0; i < N; i++ {
			wg.Add(1)
			go func(i int) {
				defer wg.Done()
				rep := make(chan batchReply, 1)
				b.submit <- &batchRequest{tag: string(rune('a' + i)), reply: rep}
				got[i] = (<-rep).json
			}(i)
		}
		wg.Wait()

		mu.Lock()
		defer mu.Unlock()
		total, maxBatch := 0, 0
		for _, s := range sizes {
			total += s
			if s > maxBatch {
				maxBatch = s
			}
		}
		Expect(total).To(Equal(N))
		Expect(maxBatch).To(BeNumerically(">=", 2), "expected at least one batch to coalesce >1 request")
	})

	It("dispatches when max size is reached", func() {
		dispatched := make(chan int, 8)
		run := func(reqs []*batchRequest) {
			dispatched <- len(reqs)
			echoReply(reqs)
		}
		b := newBatcher(2, time.Hour, run) // huge window: only size can trigger
		stop := make(chan struct{})
		go b.run(stop)
		defer close(stop)
		for i := 0; i < 2; i++ {
			rep := make(chan batchReply, 1)
			b.submit <- &batchRequest{tag: "x", reply: rep}
			go func(rep chan batchReply) { <-rep }(rep)
		}
		Eventually(dispatched, "2s").Should(Receive(Equal(2)))
	})

	It("dispatches when the wait window elapses", func() {
		dispatched := make(chan int, 8)
		run := func(reqs []*batchRequest) {
			dispatched <- len(reqs)
			echoReply(reqs)
		}
		b := newBatcher(8, 20*time.Millisecond, run) // size unreachable; window fires
		stop := make(chan struct{})
		go b.run(stop)
		defer close(stop)
		rep := make(chan batchReply, 1)
		b.submit <- &batchRequest{tag: "x", reply: rep}
		go func() { <-rep }()
		Eventually(dispatched, "2s").Should(Receive(Equal(1)))
	})

	It("bypasses batching when max size is 1", func() {
		dispatched := make(chan int, 8)
		run := func(reqs []*batchRequest) {
			dispatched <- len(reqs)
			echoReply(reqs)
		}
		b := newBatcher(1, time.Hour, run) // size 1 => immediate dispatch
		stop := make(chan struct{})
		go b.run(stop)
		defer close(stop)
		rep := make(chan batchReply, 1)
		b.submit <- &batchRequest{tag: "x", reply: rep}
		go func() { <-rep }()
		Eventually(dispatched, "2s").Should(Receive(Equal(1)))
	})
})