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

package service

import (
	"context"
	"errors"
	"testing"
	"time"

	edgenode "iop/apps/edge/internal/node"
)

// longSlotCount reads the group's long in-flight counter for a provider slot.
func longSlotCount(m *modelQueueManager, groupKey, nodeID, providerID string) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	g, ok := m.groups[groupKey]
	if !ok {
		return -1
	}
	return g.longInflight[nodeID+":"+providerID]
}

func inflightSlotCount(m *modelQueueManager, groupKey, nodeID, providerID string) int {
	m.mu.Lock()
	defer m.mu.Unlock()
	g, ok := m.groups[groupKey]
	if !ok {
		return -1
	}
	return g.inflight[nodeID+":"+providerID]
}

// TestModelQueueLongDispatchReservesBothCounters verifies S04: a long request
// occupies both the normal in-flight counter and the long in-flight counter, and
// every terminal reason (complete/error/cancelled) recovers both.
func TestModelQueueLongDispatchReservesBothCounters(t *testing.T) {
	for _, reason := range []string{"complete", "error", "cancelled"} {
		t.Run(reason, func(t *testing.T) {
			entry := &edgenode.NodeEntry{NodeID: "node-long"}
			cands := []candidateNode{{entry: entry, capacity: 3, providerID: "prov-long", longContextCapacity: 2}}
			m := newModelQueueManager(nil)

			sel, err := m.admit(context.Background(), "g-long", "", "", cands, groupPolicy{}, true)
			if err != nil || sel == nil {
				t.Fatalf("long admit: %v", err)
			}
			if got := inflightSlotCount(m, "g-long", "node-long", "prov-long"); got != 1 {
				t.Errorf("normal in-flight after admit = %d, want 1", got)
			}
			if got := longSlotCount(m, "g-long", "node-long", "prov-long"); got != 1 {
				t.Errorf("long in-flight after admit = %d, want 1", got)
			}

			longReserved := sel.longContextCapacity > 0
			m.trackInflight("g-long", "run-long-1", sel.entry.NodeID, sel.providerID, longReserved)
			if got := m.longInFlightForProvider("node-long", "prov-long"); got != 1 {
				t.Errorf("longInFlightForProvider after track = %d, want 1", got)
			}

			m.releaseRun("run-long-1", reason)
			if got := inflightSlotCount(m, "g-long", "node-long", "prov-long"); got != 0 {
				t.Errorf("normal in-flight after %s = %d, want 0", reason, got)
			}
			if got := longSlotCount(m, "g-long", "node-long", "prov-long"); got != 0 {
				t.Errorf("long in-flight after %s = %d, want 0", reason, got)
			}
			if got := m.longInFlightForProvider("node-long", "prov-long"); got != 0 {
				t.Errorf("longInFlightForProvider after %s = %d, want 0", reason, got)
			}
		})
	}
}

// TestQueueSkipsLongHeadWhenNormalBehindCanDispatch verifies the queue-skip
// contract (S06 / queue-skip): when the head of the queue is a long request
// whose provider long slots are full, a normal request queued behind it is
// still dispatched on ordinary capacity, and the blocked long head stays queued.
func TestQueueSkipsLongHeadWhenNormalBehindCanDispatch(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-skip"}
	// Provider has normal capacity 2 and a single long slot.
	cands := []candidateNode{{entry: entry, capacity: 2, providerID: "prov-skip", longContextCapacity: 1}}
	m := newModelQueueManager(nil)

	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-skip", groupPolicy{})
	// One long request already in flight: the single long slot is full while one
	// of the two normal slots remains free.
	g.inflight["node-skip:prov-skip"] = 1
	g.longInflight["node-skip:prov-skip"] = 1

	longItem := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
		long:       true,
	}
	normalItem := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	// long head first, normal request queued behind it.
	g.queue = []*queueItem{longItem, normalItem}
	m.tryDispatchLocked(g)
	m.mu.Unlock()

	// The normal request behind the blocked long head must be dispatched.
	select {
	case res := <-normalItem.waitCh:
		if res.err != nil {
			t.Fatalf("normal item expected dispatch, got error: %v", res.err)
		}
		if res.candidate == nil || res.candidate.providerID != "prov-skip" {
			t.Fatalf("normal item: unexpected candidate %v", res.candidate)
		}
	default:
		t.Fatal("normal request behind long head was not dispatched (head-of-line blocking)")
	}

	// The long head must remain queued because its long slot is still full.
	select {
	case <-longItem.waitCh:
		t.Fatal("long head should not dispatch while its long slot is full")
	default:
	}

	m.mu.Lock()
	if len(g.queue) != 1 || g.queue[0] != longItem {
		n := len(g.queue)
		m.mu.Unlock()
		t.Fatalf("expected only the long head to remain queued, got %d item(s)", n)
	}
	m.mu.Unlock()
}

// TestQueueSkipPreservesFIFOAmongDispatchableNormalItems verifies that the
// front-scan dispatch keeps FIFO order among equally-dispatchable items: with a
// single normal slot, the earlier queued item is dispatched first and the later
// one only after another slot frees.
func TestQueueSkipPreservesFIFOAmongDispatchableNormalItems(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-fifo2"}
	cands := []candidateNode{{entry: entry, capacity: 1, providerID: "prov-fifo2"}}
	m := newModelQueueManager(nil)

	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-fifo2", groupPolicy{})
	// Fill the single normal slot so both queued items must wait.
	g.inflight["node-fifo2:prov-fifo2"] = 1
	first := &queueItem{candidates: cands, waitCh: make(chan admitResult, 1), deadline: time.Now().Add(2 * time.Second)}
	second := &queueItem{candidates: cands, waitCh: make(chan admitResult, 1), deadline: time.Now().Add(2 * time.Second)}
	g.queue = []*queueItem{first, second}
	m.mu.Unlock()

	// Release one slot: only the earlier (first) item may dispatch.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo2", "node-fifo2", "prov-fifo2", false)
	m.mu.Unlock()

	select {
	case res := <-first.waitCh:
		if res.err != nil {
			t.Fatalf("first item dispatch error: %v", res.err)
		}
	default:
		t.Fatal("first (earliest) item was not dispatched first")
	}
	select {
	case <-second.waitCh:
		t.Fatal("second item should still wait (only one slot freed)")
	default:
	}

	// Release again → the second item dispatches.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo2", "node-fifo2", "prov-fifo2", false)
	m.mu.Unlock()
	select {
	case res := <-second.waitCh:
		if res.err != nil {
			t.Fatalf("second item dispatch error: %v", res.err)
		}
	default:
		t.Fatal("second item not dispatched after second slot release")
	}
}

// TestModelQueueNormalRequestDoesNotReserveLongSlot verifies a normal request
// occupies only the normal in-flight counter, never the long counter, even on a
// provider that declares long capacity.
func TestModelQueueNormalRequestDoesNotReserveLongSlot(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-n"}
	cands := []candidateNode{{entry: entry, capacity: 3, providerID: "prov-n", longContextCapacity: 2}}
	m := newModelQueueManager(nil)

	sel, err := m.admit(context.Background(), "g-n", "", "", cands, groupPolicy{}) // long omitted → normal
	if err != nil || sel == nil {
		t.Fatalf("normal admit: %v", err)
	}
	if got := inflightSlotCount(m, "g-n", "node-n", "prov-n"); got != 1 {
		t.Errorf("normal in-flight = %d, want 1", got)
	}
	if got := longSlotCount(m, "g-n", "node-n", "prov-n"); got != 0 {
		t.Errorf("long in-flight for normal request = %d, want 0", got)
	}
}

// TestModelQueueLongSlotFullExcludesLongKeepsNormal verifies S05 and D06: a
// provider whose long slots are full is excluded from long candidates (the long
// request moves to another provider or queues), but the same provider remains a
// valid candidate for normal requests while its normal capacity has room.
func TestModelQueueLongSlotFullExcludesLongKeepsNormal(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-a"}
	entryB := &edgenode.NodeEntry{NodeID: "node-b"}
	cands := []candidateNode{
		{entry: entryA, capacity: 3, priority: 0, providerID: "prov-a", longContextCapacity: 1},
		{entry: entryB, capacity: 3, priority: 1, providerID: "prov-b", longContextCapacity: 1},
	}
	m := newModelQueueManager(nil)

	// Long #1 → prov-a (priority 0 preferred), filling prov-a's single long slot.
	sel1, err := m.admit(context.Background(), "g-x", "", "", cands, groupPolicy{}, true)
	if err != nil || sel1 == nil {
		t.Fatalf("long1 admit: %v", err)
	}
	if sel1.providerID != "prov-a" {
		t.Fatalf("long1 provider = %q, want prov-a", sel1.providerID)
	}

	// Long #2 must skip the long-full prov-a and land on prov-b.
	sel2, err := m.admit(context.Background(), "g-x", "", "", cands, groupPolicy{}, true)
	if err != nil || sel2 == nil {
		t.Fatalf("long2 admit: %v", err)
	}
	if sel2.providerID != "prov-b" {
		t.Fatalf("long2 provider = %q, want prov-b (prov-a long slot full)", sel2.providerID)
	}

	// A normal request restricted to the long-full prov-a is still admitted,
	// because normal admission ignores long slot state (D06).
	normalCands := []candidateNode{cands[0]}
	selN, err := m.admit(context.Background(), "g-x", "", "", normalCands, groupPolicy{})
	if err != nil || selN == nil {
		t.Fatalf("normal admit on long-full provider: %v", err)
	}
	if selN.providerID != "prov-a" {
		t.Fatalf("normal provider = %q, want prov-a", selN.providerID)
	}

	// A further long request restricted to the long-full prov-a cannot be
	// admitted immediately; it queues and times out.
	longOnlyA := []candidateNode{cands[0]}
	_, err = m.admit(context.Background(), "g-x", "", "", longOnlyA, groupPolicy{maxQueue: 1, queueTimeout: 20 * time.Millisecond}, true)
	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("long request on long-full provider: expected queue timeout, got %v", err)
	}
}

func TestModelQueueReasonLongContextCapacityFull(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-a"}
	cands := []candidateNode{
		{entry: entryA, capacity: 3, priority: 0, providerID: "prov-a", longContextCapacity: 1},
	}
	m := newModelQueueManager(nil)

	// Long #1 → prov-a, filling prov-a's single long slot.
	_, _, err := m.admitWithReason(context.Background(), "g-x", "", "", cands, groupPolicy{}, true)
	if err != nil {
		t.Fatalf("first admit: %v", err)
	}

	// Long #2 → prov-a. It should block because longContextCapacity is full.
	done := make(chan struct{})
	go func() {
		defer close(done)
		_, _, err := m.admitWithReason(context.Background(), "g-x", "", "", cands, groupPolicy{maxQueue: 1, queueTimeout: 50 * time.Millisecond}, true)
		if !errors.Is(err, errQueueTimeout) {
			t.Errorf("expected timeout, got %v", err)
		}
	}()

	// Wait a bit for goroutine to queue the item.
	time.Sleep(10 * time.Millisecond)

	m.mu.Lock()
	g, ok := m.groups["g-x"]
	if !ok || len(g.queue) != 1 {
		m.mu.Unlock()
		t.Fatalf("expected 1 queued item, got group ok=%t, len=%d", ok, len(g.queue))
	}
	reason := g.queue[0].reason
	m.mu.Unlock()

	if reason != "long_context_capacity_full" {
		t.Errorf("expected queue reason 'long_context_capacity_full', got %q", reason)
	}

	<-done
}

func TestModelQueueReasonCapacityFull(t *testing.T) {
	entryA := &edgenode.NodeEntry{NodeID: "node-a"}
	cands := []candidateNode{
		{entry: entryA, capacity: 1, priority: 0, providerID: "prov-a"},
	}
	m := newModelQueueManager(nil)

	// #1 → prov-a, filling capacity.
	_, _, err := m.admitWithReason(context.Background(), "g-y", "", "", cands, groupPolicy{}, false)
	if err != nil {
		t.Fatalf("first admit: %v", err)
	}

	// #2 → prov-a. It should block because capacity is full.
	done := make(chan struct{})
	go func() {
		defer close(done)
		_, _, err := m.admitWithReason(context.Background(), "g-y", "", "", cands, groupPolicy{maxQueue: 1, queueTimeout: 50 * time.Millisecond}, false)
		if !errors.Is(err, errQueueTimeout) {
			t.Errorf("expected timeout, got %v", err)
		}
	}()

	time.Sleep(10 * time.Millisecond)

	m.mu.Lock()
	g, ok := m.groups["g-y"]
	if !ok || len(g.queue) != 1 {
		m.mu.Unlock()
		t.Fatalf("expected 1 queued item")
	}
	reason := g.queue[0].reason
	m.mu.Unlock()

	if reason != "capacity_full" {
		t.Errorf("expected queue reason 'capacity_full', got %q", reason)
	}

	<-done
}