notify/inbound_dispatcher_test.go

package notify

import (
	"sync"
	"sync/atomic"
	"testing"
	"time"
)

func TestInboundDispatcherSerializesPerSource(t *testing.T) {
	dispatcher := newInboundDispatcher()
	defer dispatcher.CloseAndWait()

	firstStarted := make(chan struct{}, 1)
	secondStarted := make(chan struct{}, 1)
	otherStarted := make(chan struct{}, 1)
	releaseFirst := make(chan struct{})

	var mu sync.Mutex
	var order []string

	record := func(step string) {
		mu.Lock()
		order = append(order, step)
		mu.Unlock()
	}

	if !dispatcher.Dispatch("alpha", func() {
		record("alpha-1-start")
		firstStarted <- struct{}{}
		<-releaseFirst
		record("alpha-1-end")
	}) {
		t.Fatal("dispatch alpha-1 failed")
	}
	if !dispatcher.Dispatch("alpha", func() {
		record("alpha-2-start")
		secondStarted <- struct{}{}
		record("alpha-2-end")
	}) {
		t.Fatal("dispatch alpha-2 failed")
	}
	if !dispatcher.Dispatch("beta", func() {
		record("beta-1-start")
		otherStarted <- struct{}{}
		record("beta-1-end")
	}) {
		t.Fatal("dispatch beta-1 failed")
	}

	select {
	case <-firstStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for alpha-1")
	}
	select {
	case <-otherStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for beta-1")
	}
	select {
	case <-secondStarted:
		t.Fatal("alpha-2 started before alpha-1 finished")
	case <-time.After(100 * time.Millisecond):
	}

	close(releaseFirst)

	select {
	case <-secondStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for alpha-2")
	}

	dispatcher.CloseAndWait()

	mu.Lock()
	defer mu.Unlock()
	if len(order) == 0 {
		t.Fatal("dispatch order is empty")
	}
	alpha1Start := indexOfString(order, "alpha-1-start")
	alpha1End := indexOfString(order, "alpha-1-end")
	alpha2Start := indexOfString(order, "alpha-2-start")
	beta1Start := indexOfString(order, "beta-1-start")
	if alpha1Start < 0 || alpha1End < 0 || alpha2Start < 0 || beta1Start < 0 {
		t.Fatalf("unexpected order trace: %v", order)
	}
	if alpha2Start < alpha1End {
		t.Fatalf("alpha source was not serialized: %v", order)
	}
	if beta1Start > alpha1End {
		t.Fatalf("beta source did not run in parallel window: %v", order)
	}
}

func indexOfString(list []string, target string) int {
	for idx, item := range list {
		if item == target {
			return idx
		}
	}
	return -1
}

func TestInboundDispatcherBoundsQueuedItems(t *testing.T) {
	dispatcher := newInboundDispatcherWithLimits(1, 1<<20)
	defer dispatcher.CloseAndWait()

	releaseFirst := make(chan struct{})
	firstStarted := make(chan struct{})
	if !dispatcher.DispatchSized("alpha", 1, func() {
		close(firstStarted)
		<-releaseFirst
	}) {
		t.Fatal("dispatch first item failed")
	}
	select {
	case <-firstStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for first item")
	}

	if !dispatcher.DispatchSized("alpha", 1, func() {}) {
		t.Fatal("dispatch second item failed")
	}

	thirdDone := make(chan bool, 1)
	go func() {
		thirdDone <- dispatcher.DispatchSized("alpha", 1, func() {})
	}()
	select {
	case <-thirdDone:
		t.Fatal("dispatch exceeded the queue limit without backpressure")
	case <-time.After(100 * time.Millisecond):
	}

	close(releaseFirst)
	select {
	case ok := <-thirdDone:
		if !ok {
			t.Fatal("blocked dispatch failed after room became available")
		}
	case <-time.After(time.Second):
		t.Fatal("blocked dispatch was not released when the queue drained")
	}
}

func TestInboundDispatcherBoundsQueuedBytes(t *testing.T) {
	dispatcher := newInboundDispatcherWithLimits(1000, 10)
	defer dispatcher.CloseAndWait()

	releaseFirst := make(chan struct{})
	firstStarted := make(chan struct{})
	if !dispatcher.DispatchSized("alpha", 1, func() {
		close(firstStarted)
		<-releaseFirst
	}) {
		t.Fatal("dispatch first item failed")
	}
	select {
	case <-firstStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for first item")
	}

	if !dispatcher.DispatchSized("alpha", 6, func() {}) {
		t.Fatal("dispatch second item failed")
	}

	thirdDone := make(chan bool, 1)
	go func() {
		thirdDone <- dispatcher.DispatchSized("alpha", 6, func() {})
	}()
	select {
	case <-thirdDone:
		t.Fatal("dispatch exceeded the byte limit without backpressure")
	case <-time.After(100 * time.Millisecond):
	}

	close(releaseFirst)
	select {
	case ok := <-thirdDone:
		if !ok {
			t.Fatal("blocked dispatch failed after room became available")
		}
	case <-time.After(time.Second):
		t.Fatal("blocked dispatch was not released when the queue drained")
	}
}

func TestInboundDispatcherCloseUnblocksBlockedDispatch(t *testing.T) {
	dispatcher := newInboundDispatcherWithLimits(1, 1<<20)

	releaseFirst := make(chan struct{})
	firstStarted := make(chan struct{})
	if !dispatcher.DispatchSized("alpha", 1, func() {
		close(firstStarted)
		<-releaseFirst
	}) {
		t.Fatal("dispatch first item failed")
	}
	select {
	case <-firstStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for first item")
	}
	if !dispatcher.DispatchSized("alpha", 1, func() {}) {
		t.Fatal("dispatch second item failed")
	}

	blockedDone := make(chan bool, 1)
	go func() {
		blockedDone <- dispatcher.DispatchSized("alpha", 1, func() {})
	}()
	select {
	case <-blockedDone:
		t.Fatal("dispatch exceeded the queue limit without backpressure")
	case <-time.After(100 * time.Millisecond):
	}

	closeWaitDone := make(chan struct{})
	go func() {
		dispatcher.CloseAndWait()
		close(closeWaitDone)
	}()

	select {
	case ok := <-blockedDone:
		if ok {
			t.Fatal("dispatch accepted work after the dispatcher closed")
		}
	case <-time.After(time.Second):
		t.Fatal("CloseAndWait did not release the blocked dispatch")
	}

	close(releaseFirst)
	select {
	case <-closeWaitDone:
	case <-time.After(time.Second):
		t.Fatal("CloseAndWait did not return after in-flight work finished")
	}
}

func TestInboundDispatcherConcurrentDispatchDrainsEveryItem(t *testing.T) {
	dispatcher := newInboundDispatcherWithLimits(8, 1<<20)

	const producers = 8
	const perProducer = 64
	var handled atomic.Int64
	var wg sync.WaitGroup
	for p := 0; p < producers; p++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			for i := 0; i < perProducer; i++ {
				if !dispatcher.DispatchSized("shared", 1, func() {
					handled.Add(1)
					time.Sleep(50 * time.Microsecond)
				}) {
					t.Errorf("dispatch rejected before close")
					return
				}
			}
		}()
	}
	wg.Wait()
	dispatcher.CloseAndWait()

	if got, want := handled.Load(), int64(producers*perProducer); got != want {
		t.Fatalf("handled=%d, want %d", got, want)
	}
	dispatcher.mu.Lock()
	queued, queuedBytes := dispatcher.queued, dispatcher.queuedBytes
	dispatcher.mu.Unlock()
	if queued != 0 || queuedBytes != 0 {
		t.Fatalf("queue accounting leaked: queued=%d bytes=%d", queued, queuedBytes)
	}
}

func TestInboundDispatcherPerSourceLimitDoesNotStarveOtherSources(t *testing.T) {
	dispatcher := newInboundDispatcherWithCaps(64, 1<<20, 1, 1<<20)
	defer dispatcher.CloseAndWait()

	releaseFirst := make(chan struct{})
	firstStarted := make(chan struct{})
	if !dispatcher.DispatchSized("alpha", 1, func() {
		close(firstStarted)
		<-releaseFirst
	}) {
		t.Fatal("dispatch first alpha item failed")
	}
	select {
	case <-firstStarted:
	case <-time.After(time.Second):
		t.Fatal("timed out waiting for first alpha item")
	}
	if !dispatcher.DispatchSized("alpha", 1, func() {}) {
		t.Fatal("dispatch second alpha item failed")
	}

	alphaBlocked := make(chan bool, 1)
	go func() {
		alphaBlocked <- dispatcher.DispatchSized("alpha", 1, func() {})
	}()
	select {
	case <-alphaBlocked:
		t.Fatal("alpha exceeded its per-source queue budget")
	case <-time.After(100 * time.Millisecond):
	}

	betaDone := make(chan bool, 1)
	go func() {
		betaDone <- dispatcher.DispatchSized("beta", 1, func() {})
	}()
	select {
	case ok := <-betaDone:
		if !ok {
			t.Fatal("beta dispatch failed")
		}
	case <-time.After(time.Second):
		t.Fatal("a saturated alpha source starved the beta source")
	}

	close(releaseFirst)
	select {
	case ok := <-alphaBlocked:
		if !ok {
			t.Fatal("blocked alpha dispatch failed after room became available")
		}
	case <-time.After(time.Second):
		t.Fatal("blocked alpha dispatch was not released")
	}
}