easyai-ai-gateway/apps/api/internal/runner/queue_worker_manager_test.go

package runner

import (
	"context"
	"errors"
	"io"
	"log/slog"
	"sync/atomic"
	"testing"
	"time"

	"github.com/easyai/easyai-ai-gateway/apps/api/internal/config"
	"github.com/easyai/easyai-ai-gateway/apps/api/internal/store"
)

type fakeAsyncExecutionClient struct {
	startErr      error
	stopGate      <-chan struct{}
	started       atomic.Int64
	stopped       atomic.Int64
	stopCancelled atomic.Int64
}

func (c *fakeAsyncExecutionClient) Start(context.Context) error {
	c.started.Add(1)
	return c.startErr
}

func (c *fakeAsyncExecutionClient) Stop(ctx context.Context) error {
	c.stopped.Add(1)
	if c.stopGate == nil {
		return nil
	}
	select {
	case <-c.stopGate:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}

func (c *fakeAsyncExecutionClient) StopAndCancel(context.Context) error {
	c.stopCancelled.Add(1)
	return nil
}

func TestResizeAsyncWorkerCapacityStartsReplacementBeforeGracefulDrain(t *testing.T) {
	oldClient := &fakeAsyncExecutionClient{}
	drainGate := make(chan struct{})
	oldClient.stopGate = drainGate
	newClient := &fakeAsyncExecutionClient{}
	service := asyncWorkerManagerTestService(1)
	service.riverExecutionClient = oldClient
	service.asyncCapacityLoader = fixedAsyncCapacity(3)
	service.asyncClientFactory = func(capacity int) (asyncExecutionClient, error) {
		if capacity != 3 {
			t.Fatalf("factory capacity=%d, want=3", capacity)
		}
		return newClient, nil
	}

	service.resizeAsyncWorkerCapacity(context.Background())
	if newClient.started.Load() != 1 {
		t.Fatalf("replacement start count=%d, want=1", newClient.started.Load())
	}
	if service.riverExecutionClient != newClient || service.riverWorkerCapacity != 3 {
		t.Fatalf("replacement was not installed: capacity=%d client=%T", service.riverWorkerCapacity, service.riverExecutionClient)
	}
	waitForAtomicValue(t, &oldClient.stopped, 1)
	if oldClient.stopCancelled.Load() != 0 {
		t.Fatal("graceful drain cancelled an already running old task")
	}
	close(drainGate)
	waitForDrainingClients(t, service, 0)
}

func TestResizeAsyncWorkerCapacityFailuresKeepCurrentClient(t *testing.T) {
	t.Run("refresh failure", func(t *testing.T) {
		oldClient := &fakeAsyncExecutionClient{}
		service := asyncWorkerManagerTestService(4)
		service.riverExecutionClient = oldClient
		service.asyncCapacityLoader = func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
			return store.AsyncWorkerCapacitySnapshot{}, errors.New("database unavailable")
		}
		service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
			t.Fatal("factory must not run after refresh failure")
			return nil, nil
		}
		service.resizeAsyncWorkerCapacity(context.Background())
		if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
			t.Fatal("refresh failure replaced the current client")
		}
	})

	t.Run("create failure", func(t *testing.T) {
		oldClient := &fakeAsyncExecutionClient{}
		service := asyncWorkerManagerTestService(4)
		service.riverExecutionClient = oldClient
		service.asyncCapacityLoader = fixedAsyncCapacity(8)
		service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
			return nil, errors.New("factory failed")
		}
		service.resizeAsyncWorkerCapacity(context.Background())
		if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
			t.Fatal("create failure replaced the current client")
		}
	})

	t.Run("start failure", func(t *testing.T) {
		oldClient := &fakeAsyncExecutionClient{}
		failedClient := &fakeAsyncExecutionClient{startErr: errors.New("start failed")}
		service := asyncWorkerManagerTestService(4)
		service.riverExecutionClient = oldClient
		service.asyncCapacityLoader = fixedAsyncCapacity(8)
		service.asyncClientFactory = func(int) (asyncExecutionClient, error) {
			return failedClient, nil
		}
		service.resizeAsyncWorkerCapacity(context.Background())
		if service.riverExecutionClient != oldClient || service.riverWorkerCapacity != 4 {
			t.Fatal("start failure replaced the current client")
		}
		if failedClient.stopCancelled.Load() != 1 {
			t.Fatal("failed replacement client was not cleaned up")
		}
	})
}

func TestStopAsyncWorkersIncludesCurrentAndDrainingClients(t *testing.T) {
	current := &fakeAsyncExecutionClient{}
	draining := &fakeAsyncExecutionClient{}
	service := asyncWorkerManagerTestService(2)
	service.riverExecutionClient = current
	service.riverDrainingClients[draining] = struct{}{}
	ctx, cancel := context.WithCancel(context.Background())
	cancel()
	service.stopAsyncWorkersOnShutdown(ctx)
	if current.stopCancelled.Load() != 1 || draining.stopCancelled.Load() != 1 {
		t.Fatalf("shutdown cancellations current=%d draining=%d, want 1/1", current.stopCancelled.Load(), draining.stopCancelled.Load())
	}
}

func asyncWorkerManagerTestService(capacity int) *Service {
	return &Service{
		cfg: config.Config{
			AsyncWorkerHardLimit:              2048,
			AsyncWorkerRefreshIntervalSeconds: 5,
		},
		logger:               slog.New(slog.NewTextHandler(io.Discard, nil)),
		riverDrainingClients: make(map[asyncExecutionClient]struct{}),
		riverWorkerCapacity:  capacity,
		riverExecutionClient: nil,
	}
}

func fixedAsyncCapacity(capacity int) func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
	return func(context.Context, int) (store.AsyncWorkerCapacitySnapshot, error) {
		return store.AsyncWorkerCapacitySnapshot{Capacity: capacity, Desired: capacity, HardLimit: 2048}, nil
	}
}

func waitForAtomicValue(t *testing.T, value *atomic.Int64, want int64) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		if value.Load() == want {
			return
		}
		time.Sleep(time.Millisecond)
	}
	t.Fatalf("atomic value=%d, want=%d", value.Load(), want)
}

func waitForDrainingClients(t *testing.T, service *Service, want int) {
	t.Helper()
	deadline := time.Now().Add(time.Second)
	for time.Now().Before(deadline) {
		service.riverMu.RLock()
		count := len(service.riverDrainingClients)
		service.riverMu.RUnlock()
		if count == want {
			return
		}
		time.Sleep(time.Millisecond)
	}
	service.riverMu.RLock()
	count := len(service.riverDrainingClients)
	service.riverMu.RUnlock()
	t.Fatalf("draining client count=%d, want=%d", count, want)
}