iop/apps/edge/internal/service/provider_pool_admission_test.go

package service

import (
	"context"
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
)

// requireProviderPoolPending polls until the manager's provider-pool pending
// count reaches at least want, or times out. Used by tests that must observe
// a specific number of queued provider-pool items before submitting the next
// admission that should be rejected by the global cap.
func requireProviderPoolPending(t *testing.T, m *modelQueueManager, want int) {
	t.Helper()
	deadline := time.Now().Add(500 * time.Millisecond)
	got := 0
	for time.Now().Before(deadline) {
		m.mu.Lock()
		got = m.pendingProviderPoolCountLocked()
		m.mu.Unlock()
		if got >= want {
			return
		}
		time.Sleep(1 * time.Millisecond)
	}
	t.Fatalf("expected provider-pool pending >= %d, reached %d (timeout)", want, got)
}

// providerPoolAdmitResult captures the outcome of a provider-pool admit call
// run in a goroutine so the test can observe dispatch vs. queue vs. rejection.
type providerPoolAdmitResult struct {
	candidate *candidateNode
	err       error
}

// TestProviderPoolMaxQueueCountsPendingAcrossGroups verifies that the global
// max_queue bound applies across all model groups for provider-pool admissions,
// that legacy adapter pending items are excluded from the count, and that the
// rejection is immediate (no fallback to group default timeout).
func TestProviderPoolMaxQueueCountsPendingAcrossGroups(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pp1",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: "prov-A", Capacity: 1},
		},
	})
	entry := &edgenode.NodeEntry{NodeID: "node-pp1"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-A"}}

	m := newModelQueueManager(store)
	m.setProviderPoolPolicyLocked(store, NewGroupPolicy(2, 30*time.Second))

	// Dispatch one provider-pool item (fills the only slot).
	first, err := m.admit(t.Context(), "group-A", "", "", cands, groupPolicy{}, nil, false, true)
	if err != nil || first == nil {
		t.Fatalf("group-A admit: %v", err)
	}

	// Submit two more items concurrently. The provider is full, so they enter
	// the queue as pending items. Both must be queued before the cap check
	// rejects the third. Use a cancellable context so the goroutines can be
	// cleaned up after the test completes.
	const ppGoroutines = 2
	resultCh := make(chan providerPoolAdmitResult, ppGoroutines)
	var wg sync.WaitGroup
	ppCtx, ppCancel := context.WithCancel(t.Context())
	defer ppCancel()
	for i := 0; i < ppGoroutines; i++ {
		wg.Add(1)
		idx := i
		go func() {
			defer wg.Done()
			c, err := m.admit(ppCtx, fmt.Sprintf("group-pp-%d", idx), "", "", cands, groupPolicy{}, nil, false, true)
			resultCh <- providerPoolAdmitResult{candidate: c, err: err}
		}()
	}

	// Wait for exactly 2 provider-pool items to be pending. Both goroutines
	// must have entered the queue before the cap blocks the next admit.
	requireProviderPoolPending(t, m, 2)

	// The fourth admission must be rejected by the global cap immediately
	// with errQueueFull, without waiting for the group default 30s timeout.
	ctx4, cancel4 := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer cancel4()
	started := time.Now()
	_, err4 := m.admit(ctx4, "group-C", "", "", cands, groupPolicy{}, nil, false, true)
	elapsed := time.Since(started)
	if !errors.Is(err4, errQueueFull) {
		t.Fatalf("expected immediate global errQueueFull, got %v after %v", err4, elapsed)
	}
	if elapsed > 200*time.Millisecond {
		t.Fatalf("errQueueFull was not immediate (expected <200ms): %v", elapsed)
	}

	// Legacy adapter admission must not be counted against the global cap:
	// the legacy path uses the per-group queue, which is independent of the
	// provider-pool policy. With concurrency=1 on the node and a legacy
	// candidate that does not carry a providerID, the legacy admission
	// creates its own group-local slot.
	legacyEntry := &edgenode.NodeEntry{NodeID: "node-pp1"}
	_, err = m.admit(t.Context(), "group-legacy", "", "", []candidateNode{{entry: legacyEntry, capacity: 1}}, groupPolicy{}, nil, false, false)
	if err != nil {
		t.Fatalf("legacy admit should not be affected by global cap: %v", err)
	}

	// Verify the pending count is still exactly 2 (legacy is excluded).
	requireProviderPoolPending(t, m, 2)

	// Cancel the pending goroutines so they exit via ctx.Done().
	ppCancel()
	wg.Wait()
	close(resultCh)
	for range resultCh {
	}
}

// TestProviderPoolCommonQueueTimeoutAcrossCandidateOrders verifies that the
// common timeout from the root policy applies regardless of candidate order.
// Two providers with equal capacity are presented in reversed order; both
// slots are filled so the waiter must queue in either ordering, and the
// timeout and error identity must be identical.
func TestProviderPoolCommonQueueTimeoutAcrossCandidateOrders(t *testing.T) {
	rootTimeout := 60 * time.Millisecond

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ord",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: "prov-A", Capacity: 1},
			{ID: "prov-B", Capacity: 1},
		},
	})
	entryA := &edgenode.NodeEntry{NodeID: "node-ord"}
	entryB := &edgenode.NodeEntry{NodeID: "node-ord"}

	orders := []struct {
		name       string
		candidates []candidateNode
	}{
		{
			name:       "A_then_B",
			candidates: []candidateNode{{entry: entryA, capacity: 1, providerID: "prov-A"}, {entry: entryB, capacity: 1, providerID: "prov-B"}},
		},
		{
			name:       "B_then_A",
			candidates: []candidateNode{{entry: entryB, capacity: 1, providerID: "prov-B"}, {entry: entryA, capacity: 1, providerID: "prov-A"}},
		},
	}

	for _, tc := range orders {
		t.Run(tc.name, func(t *testing.T) {
			m := newModelQueueManager(store)
			m.setProviderPoolPolicyLocked(store, NewGroupPolicy(16, rootTimeout))

			// Fill both slots so the waiter must queue regardless of candidate order.
			_, err := m.admit(t.Context(), "group-ord-init-1", "", "", tc.candidates, groupPolicy{}, nil, false, true)
			if err != nil {
				t.Fatalf("initial admit 1: %v", err)
			}
			_, err = m.admit(t.Context(), "group-ord-init-2", "", "", tc.candidates, groupPolicy{}, nil, false, true)
			if err != nil {
				t.Fatalf("initial admit 2: %v", err)
			}

			// The waiter uses the same candidate slice; root policy timeout
			// must govern regardless of candidate order.
			start := time.Now()
			ctx, cancel := context.WithTimeout(t.Context(), 500*time.Millisecond)
			defer cancel()
			_, err = m.admit(ctx, "group-ord-waiter", "", "", tc.candidates, groupPolicy{}, nil, false, true)
			elapsed := time.Since(start)

			if err == nil {
				t.Fatal("expected timeout error, got nil")
			}
			if !errors.Is(err, errQueueTimeout) {
				t.Fatalf("expected errQueueTimeout, got %v", err)
			}
			// Root timeout of 60ms is intentionally tight so a regression
			// that reverts to the group default 30s is caught. The tolerance
			// allows for goroutine-scheduler variance under load (a slow CI
			// can easily delay timer expiry by 100+ms, and the -race build
			// adds further overhead). Anything beyond 5x the root timeout is
			// a real drift, so fail hard there.
			if elapsed > rootTimeout*5 {
				t.Errorf("timeout took %v (expected ~%v, within 5x margin)", elapsed, rootTimeout)
			}
		})
	}
}

// TestProviderPoolZeroQueueTimeoutWaitsForContext verifies that explicit zero
// queue_timeout_ms preserves no-timeout semantics: the internal timer is not
// created, the waiter only returns on ctx.Done(), and the error identity is
// ctx.Err() (context.DeadlineExceeded), never errQueueTimeout.
func TestProviderPoolZeroQueueTimeoutWaitsForContext(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ppz",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: "prov-Z", Capacity: 1},
		},
	})
	entry := &edgenode.NodeEntry{NodeID: "node-ppz"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-Z"}}

	m := newModelQueueManager(store)
	policy := NewGroupPolicy(16, 0)
	m.setProviderPoolPolicyLocked(store, policy)

	// Verify the constructor preserved the explicit zero.
	if !policy.queueTimeoutSet {
		t.Fatalf("NewGroupPolicy(max, 0) must set queueTimeoutSet=true, got %v", policy.queueTimeoutSet)
	}
	if policy.queueTimeout != 0 {
		t.Fatalf("NewGroupPolicy(max, 0) must preserve queueTimeout=0, got %v", policy.queueTimeout)
	}

	// Fill the slot.
	_, err := m.admit(t.Context(), "group-z-1", "", "", cands, groupPolicy{}, nil, false, true)
	if err != nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Waiter with a 80ms context. With explicit zero timeout the internal
	// timer must not be set, so the only exit is ctx.Done().
	ctx, cancel := context.WithTimeout(t.Context(), 80*time.Millisecond)
	defer cancel()
	start := time.Now()
	_, err = m.admit(ctx, "group-z-2", "", "", cands, groupPolicy{}, nil, false, true)
	elapsed := time.Since(start)

	if err == nil {
		t.Fatal("expected context cancellation, got nil")
	}
	if !errors.Is(err, context.DeadlineExceeded) {
		t.Fatalf("expected context.DeadlineExceeded (no internal timer), got %v", err)
	}
	if strings.Contains(err.Error(), "queue timeout") {
		t.Fatalf("zero timeout must not produce errQueueTimeout: %v", err)
	}
	// Must complete close to the context deadline, not the group default 30s.
	if elapsed > 250*time.Millisecond {
		t.Fatalf("waiter did not return after context cancellation: %v", elapsed)
	}
}

// TestProviderPoolPendingCountRecoversAfterCancelAndTimeout verifies that the
// global pending count, lease table, and provider resource state all recover
// after a mix of explicit cancel and independent queue timeout among pending
// waiters. Error identity (context.Canceled vs. errQueueTimeout) is verified
// per waiter to confirm each path runs independently.
func TestProviderPoolPendingCountRecoversAfterCancelAndTimeout(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pp4",
		Runtime: config.RuntimeConf{Concurrency: 1},
		Providers: []config.NodeProviderConf{
			{ID: "prov-W", Capacity: 1},
		},
	})
	entry := &edgenode.NodeEntry{NodeID: "node-pp4"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-W"}}

	m := newModelQueueManager(store)
	// Use 1s queue timeout so the second waiter's independent internal timeout
	// is not affected by race instrumentation overhead that previously caused
	// both waiters to miss the 80ms barrier under load.
	m.setProviderPoolPolicyLocked(store, NewGroupPolicy(2, time.Second))

	// Fill the slot.
	_, err := m.admit(t.Context(), "group-1", "", "", cands, groupPolicy{}, nil, false, true)
	if err != nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Staged enqueue: start the first waiter, observe it queued, then start
	// the second. This eliminates the timing-dependent assumption that both
	// goroutines enter the queue within 80ms simultaneously.
	ch1 := make(chan error, 1)
	ctx1, cancel1 := context.WithCancel(t.Context())
	go func() {
		_, err := m.admit(ctx1, "group-2", "", "", cands, groupPolicy{}, nil, false, true)
		ch1 <- err
	}()
	requireProviderPoolPending(t, m, 1)

	ch2 := make(chan error, 1)
	ctx2, cancel2 := context.WithTimeout(t.Context(), time.Second+500*time.Millisecond)
	defer cancel2()
	go func() {
		_, err := m.admit(ctx2, "group-3", "", "", cands, groupPolicy{}, nil, false, true)
		ch2 <- err
	}()
	requireProviderPoolPending(t, m, 2)

	// Cancel waiter 1 explicitly — it must exit via context.Canceled, not
	// through the queue manager's internal timer.
	cancel1()

	var err1 error
	select {
	case err1 = <-ch1:
	case <-time.After(1 * time.Second):
		t.Fatal("waiter 1 did not return after cancel")
	}
	if err1 == nil {
		t.Fatal("waiter 1 expected error, got nil")
	}
	if !errors.Is(err1, context.Canceled) {
		t.Fatalf("waiter 1 expected context.Canceled, got %v", err1)
	}

	// Waiter 2 stays queued until the internal 1s queue timeout fires.
	// The independent internal timeout is the determining exit path.
	var err2 error
	select {
	case err2 = <-ch2:
	case <-time.After(2 * time.Second):
		t.Fatal("waiter 2 did not return within timeout")
	}
	if err2 == nil {
		t.Fatal("waiter 2 expected error, got nil")
	}
	if !errors.Is(err2, errQueueTimeout) {
		t.Fatalf("waiter 2 expected errQueueTimeout, got %v", err2)
	}

	// Pending count must have recovered to 0.
	m.mu.Lock()
	pending := m.pendingProviderPoolCountLocked()
	m.mu.Unlock()
	if pending != 0 {
		t.Fatalf("expected pending count=0 after cancel+timeout, got %d", pending)
	}

	// Lease table must have exactly 1 lease (the initial dispatch).
	if lc := leaseCount(m); lc != 1 {
		t.Fatalf("expected leaseCount=1 (initial dispatch only), got %d", lc)
	}

	// Provider resource in-flight must not have changed from the initial admit.
	inFlight, _ := providerResourceCounts(m, "node-pp4", "prov-W")
	if inFlight != 1 {
		t.Fatalf("expected prov-W in-flight=1 (only initial dispatch), got %d", inFlight)
	}
}

// TestProviderPoolMaxQueueIgnoresLegacyPending verifies that legacy adapter
// pending items are not counted against the provider-pool global cap. A
// provider-pool item is first dispatched so the provider resource fills to
// capacity, then legacy items dispatch and a legacy waiter enters the group
// queue. Separately, a second provider-pool waiter is rejected by the global
// cap even though a legacy pending item exists, because legacy pending is
// excluded from the provider-pool count. Explicit cancel must yield
// context.Canceled for both waiters, and releasing all holder leases must
// recover pending, lease, legacy group in-flight, and provider resource
// in-flight to zero.
func TestProviderPoolMaxQueueIgnoresLegacyPending(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-pp5",
		Runtime: config.RuntimeConf{Concurrency: 2},
		Providers: []config.NodeProviderConf{
			{ID: "prov-Legacy", Capacity: 1},
		},
	})
	entry := &edgenode.NodeEntry{NodeID: "node-pp5"}

	// Legacy candidates: NO providerID — matches the production legacy
	// interpretation where providerID is only assigned in the provider-pool
	// resolution path (resolveProviderPoolCandidates). Legacy admission goes
	// through group-local inflight accounting (slotKey = nodeID only), which
	// is independent of the provider resource tracked in m.resources.
	legacyCandidates := []candidateNode{{entry: entry, capacity: 1}}

	// Provider-pool candidates: WITH providerID — matches the production
	// provider-pool interpretation where resolveProviderPoolCandidates fills
	// in the providerID from the catalog-matched provider.
	providerCandidates := []candidateNode{{
		entry: entry, capacity: 1, providerID: "prov-Legacy",
	}}

	m := newModelQueueManager(store)
	// max_queue=1: exactly one slot at the global cap. Legacy pending is
	// excluded from this count, so a provider-pool waiter can enter the
	// queue even though a legacy item is pending. The next provider-pool
	// admission must then be rejected by errQueueFull immediately.
	m.setProviderPoolPolicyLocked(store, NewGroupPolicy(1, 30*time.Second))

	// === Provider-pool phase 1: dispatch 1 item to fill the prov resource ===
	// This ensures subsequent provider-pool admissions find the resource at
	// capacity and must queue against the global cap.
	var ppLeaseIDs []uint64
	c, err := m.admit(t.Context(), "pp-group", "", "", providerCandidates, groupPolicy{}, nil, false, true)
	if err != nil || c == nil {
		t.Fatalf("pp admit 1: %v", err)
	}
	ppLeaseIDs = append(ppLeaseIDs, c.leaseID)

	// === Legacy phase: dispatch 1 item (legacy inflight, separate from prov resource) ===
	// Legacy uses group-local inflight (slotKey = "node-pp5"). The dispatch
	// fills that group-local slot. This is independent of the provider
	// resource state used by provider-pool admissions.
	var legacyLeaseIDs []uint64
	c, err = m.admit(t.Context(), "lg-group", "", "", legacyCandidates, groupPolicy{}, nil, false, false)
	if err != nil || c == nil {
		t.Fatalf("legacy admit 1: %v", err)
	}
	legacyLeaseIDs = append(legacyLeaseIDs, c.leaseID)

	m.mu.Lock()
	lgInflight := m.groups["lg-group"].inflight["node-pp5"]
	m.mu.Unlock()
	if lgInflight != 1 {
		t.Fatalf("expected legacy group inflight=1 after 1 dispatched item, got %d", lgInflight)
	}

	// === Start legacy waiter (legacy group is already at capacity: inflight=1) ===
	ctxLegacy, cancelLegacy := context.WithCancel(t.Context())
	defer cancelLegacy()
	chLegacy := make(chan error, 1)
	go func() {
		_, err := m.admit(ctxLegacy, "lg-group", "", "", legacyCandidates, groupPolicy{}, nil, false, false)
		chLegacy <- err
	}()
	waitForQueueLen(t, m, "lg-group", 1)

	// === Start provider-pool waiter (prov resource full → enters queue) ===
	// Provider-pool uses m.resources[{nodeID, providerID}] accounting.
	// Legacy pending is excluded from the global cap, so the second
	// provider-pool waiter can still enter the queue.
	ctxPp, cancelPp := context.WithCancel(t.Context())
	defer cancelPp()
	chPp := make(chan error, 1)
	go func() {
		_, err := m.admit(ctxPp, "pp-waiter", "", "", providerCandidates, groupPolicy{}, nil, false, true)
		chPp <- err
	}()
	requireProviderPoolPending(t, m, 1)

	// === Cap boundary: second provider-pool admission must be rejected ===
	// This is the actual global cap boundary assertion: legacy pending exists
	// but is excluded, so one provider-pool waiter occupies the queue slot,
	// and the next provider-pool admission must return errQueueFull immediately.
	ctxCap, cancelCap := context.WithTimeout(t.Context(), 500*time.Millisecond)
	defer cancelCap()
	started := time.Now()
	_, errCap := m.admit(ctxCap, "pp-over-cap", "", "", providerCandidates, groupPolicy{}, nil, false, true)
	if !errors.Is(errCap, errQueueFull) {
		t.Fatalf("expected legacy-excluded provider-pool errQueueFull, got %v", errCap)
	}
	if elapsed := time.Since(started); elapsed > 200*time.Millisecond {
		t.Fatalf("errQueueFull was not immediate: %v", elapsed)
	}

	// === Verify isolation: legacy pending NOT counted in pp pending ===
	m.mu.Lock()
	ppPending := m.pendingProviderPoolCountLocked()
	m.mu.Unlock()
	if ppPending != 1 {
		t.Fatalf("expected pp pending=1 (only provider-pool waiter), got %d", ppPending)
	}

	// Verify legacy pending item exists in the legacy group.
	m.mu.Lock()
	lg, ok := m.groups["lg-group"]
	legacyPending := 0
	if ok {
		for _, item := range lg.queue {
			if !item.providerPool {
				legacyPending++
			}
		}
	}
	m.mu.Unlock()
	if legacyPending != 1 {
		t.Fatalf("expected 1 legacy pending item in group, got %d", legacyPending)
	}

	// === Cancel both waiters explicitly — each must exit via context.Canceled ===
	cancelLegacy()
	var errLegacy error
	select {
	case errLegacy = <-chLegacy:
	case <-time.After(1 * time.Second):
		t.Fatal("legacy waiter did not return after cancel")
	}
	if errLegacy == nil {
		t.Fatal("legacy waiter expected error, got nil")
	}
	if !errors.Is(errLegacy, context.Canceled) {
		t.Fatalf("legacy waiter expected context.Canceled, got %v", errLegacy)
	}

	cancelPp()
	var errPp error
	select {
	case errPp = <-chPp:
	case <-time.After(500 * time.Millisecond):
		t.Fatal("provider-pool waiter did not return after cancel")
	}
	if errPp == nil {
		t.Fatal("provider-pool waiter expected error, got nil")
	}
	if !errors.Is(errPp, context.Canceled) {
		t.Fatalf("provider-pool waiter expected context.Canceled, got %v", errPp)
	}

	// === Recovery: release all holder leases, verify everything returns to 0 ===
	for _, id := range ppLeaseIDs {
		m.releaseLease(id, "test-recovery")
	}
	for _, id := range legacyLeaseIDs {
		m.releaseLease(id, "test-recovery")
	}

	// Pending count must be 0.
	m.mu.Lock()
	ppPending = m.pendingProviderPoolCountLocked()
	m.mu.Unlock()
	if ppPending != 0 {
		t.Fatalf("expected pp pending=0 after holder release, got %d", ppPending)
	}

	// Lease table must be empty (all holder leases released).
	if lc := leaseCount(m); lc != 0 {
		t.Fatalf("expected leaseCount=0 after holder release, got %d", lc)
	}

	// Provider resource in-flight must be 0.
	inFlight, _ := providerResourceCounts(m, "node-pp5", "prov-Legacy")
	if inFlight != 0 {
		t.Fatalf("expected prov resource inFlight=0 after holder release, got %d", inFlight)
	}

	// Sum of all legacy group inflight slots must be 0.
	m.mu.Lock()
	totalLgInflight := 0
	for _, g := range m.groups {
		for _, v := range g.inflight {
			totalLgInflight += v
		}
	}
	m.mu.Unlock()
	if totalLgInflight != 0 {
		t.Fatalf("expected total legacy group inflight=0 after holder release, got %d", totalLgInflight)
	}
}