nnet/pkg/executor/executor_test.go

package executor

import (
	"context"
	"errors"
	"sync"
	"testing"
	"time"
)

func TestPerKeyPreservesOrderAndRunsDifferentKeysConcurrently(t *testing.T) {
	e := NewPerKey(Config{Workers: 2, QueueSize: 8})
	defer func() {
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		if err := e.Shutdown(ctx); err != nil {
			t.Fatalf("shutdown error: %v", err)
		}
	}()

	var mu sync.Mutex
	var order []int
	firstStarted := make(chan struct{})
	releaseFirst := make(chan struct{})
	secondDone := make(chan struct{})

	if err := e.Submit("a", func() {
		close(firstStarted)
		<-releaseFirst
		mu.Lock()
		order = append(order, 1)
		mu.Unlock()
	}); err != nil {
		t.Fatalf("submit first task: %v", err)
	}
	if err := e.Submit("a", func() {
		mu.Lock()
		order = append(order, 2)
		mu.Unlock()
	}); err != nil {
		t.Fatalf("submit second task: %v", err)
	}
	if err := e.Submit("b", func() {
		close(secondDone)
	}); err != nil {
		t.Fatalf("submit different-key task: %v", err)
	}

	select {
	case <-firstStarted:
	case <-time.After(time.Second):
		t.Fatal("first task did not start")
	}
	select {
	case <-secondDone:
	case <-time.After(time.Second):
		t.Fatal("different-key task did not run concurrently")
	}
	close(releaseFirst)

	deadline := time.After(time.Second)
	for {
		mu.Lock()
		done := len(order) == 2
		mu.Unlock()
		if done {
			break
		}
		select {
		case <-deadline:
			t.Fatal("same-key tasks did not drain")
		default:
			time.Sleep(time.Millisecond)
		}
	}

	mu.Lock()
	defer mu.Unlock()
	if order[0] != 1 || order[1] != 2 {
		t.Fatalf("order = %v, want [1 2]", order)
	}
}

func TestPerKeyRejectsFullQueueWithoutBlocking(t *testing.T) {
	e := NewPerKey(Config{Workers: 1, QueueSize: 1})
	defer func() {
		closeTask := make(chan struct{})
		_ = e.Submit("a", func() { close(closeTask) })
		ctx, cancel := context.WithTimeout(context.Background(), time.Second)
		defer cancel()
		_ = e.Shutdown(ctx)
	}()

	started := make(chan struct{})
	release := make(chan struct{})
	if err := e.Submit("a", func() {
		close(started)
		<-release
	}); err != nil {
		t.Fatalf("submit running task: %v", err)
	}
	select {
	case <-started:
	case <-time.After(time.Second):
		t.Fatal("running task did not start")
	}

	if err := e.Submit("a", func() {}); err != nil {
		t.Fatalf("submit queued task: %v", err)
	}

	done := make(chan error, 1)
	go func() {
		done <- e.Submit("a", func() {})
	}()
	select {
	case err := <-done:
		if !errors.Is(err, ErrQueueFull) {
			t.Fatalf("submit error = %v, want ErrQueueFull", err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("queue-full submit blocked")
	}
	close(release)
}

func TestPerKeyShutdownDrainsAcceptedTasks(t *testing.T) {
	e := NewPerKey(Config{Workers: 1, QueueSize: 4})
	var countMu sync.Mutex
	count := 0
	for i := 0; i < 3; i++ {
		if err := e.Submit("a", func() {
			countMu.Lock()
			count++
			countMu.Unlock()
		}); err != nil {
			t.Fatalf("submit task %d: %v", i, err)
		}
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := e.Shutdown(ctx); err != nil {
		t.Fatalf("shutdown error: %v", err)
	}
	if err := e.Submit("a", func() {}); !errors.Is(err, ErrClosed) {
		t.Fatalf("submit after shutdown error = %v, want ErrClosed", err)
	}

	countMu.Lock()
	defer countMu.Unlock()
	if count != 3 {
		t.Fatalf("completed tasks = %d, want 3", count)
	}
}

func TestPerKeyRecoversTaskPanic(t *testing.T) {
	panicSeen := make(chan interface{}, 1)
	e := NewPerKey(Config{
		Workers:   1,
		QueueSize: 2,
		OnPanic: func(value interface{}) {
			panicSeen <- value
		},
	})
	if err := e.Submit("a", func() { panic("boom") }); err != nil {
		t.Fatalf("submit panic task: %v", err)
	}
	if err := e.Submit("a", func() {}); err != nil {
		t.Fatalf("submit follow-up task: %v", err)
	}

	select {
	case value := <-panicSeen:
		if value != "boom" {
			t.Fatalf("panic value = %v, want boom", value)
		}
	case <-time.After(time.Second):
		t.Fatal("panic callback was not called")
	}

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()
	if err := e.Shutdown(ctx); err != nil {
		t.Fatalf("shutdown error: %v", err)
	}
}