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

package service

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

	edgeevents "iop/apps/edge/internal/events"
	edgenode "iop/apps/edge/internal/node"
	"iop/packages/go/config"
	iop "iop/proto/gen/iop"
)

// waitForQueueLen polls until the group queue reaches wantLen or times out.
func waitForQueueLen(t *testing.T, m *modelQueueManager, groupKey string, wantLen int) {
	t.Helper()
	deadline := time.Now().Add(200 * time.Millisecond)
	for time.Now().Before(deadline) {
		m.mu.Lock()
		g, ok := m.groups[groupKey]
		got := 0
		if ok {
			got = len(g.queue)
		}
		m.mu.Unlock()
		if got >= wantLen {
			return
		}
		time.Sleep(1 * time.Millisecond)
	}
	t.Fatalf("timeout: queue for %q did not reach length %d", groupKey, wantLen)
}

// TestModelQueueFIFOOrdering verifies that queued items are dispatched in FIFO
// order when a slot becomes available.
func TestModelQueueFIFOOrdering(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-q1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-q1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(store)

	// Fill the only capacity slot.
	first, err := m.admit(context.Background(), "g-fifo", cands, defPolicy)
	if err != nil || first == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Manually insert two items in known order so we can verify FIFO.
	item1 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	item2 := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-fifo", groupPolicy{})
	g.queue = append(g.queue, item1, item2)
	m.mu.Unlock()

	// Release one slot — item1 (head) must be dispatched first.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1")
	m.mu.Unlock()

	select {
	case res := <-item1.waitCh:
		if res.err != nil {
			t.Fatalf("item1 expected dispatch, got error: %v", res.err)
		}
		if res.node == nil || res.node.NodeID != "node-q1" {
			t.Fatalf("item1: unexpected node %v", res.node)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item1 was not dispatched after slot release")
	}

	// item2 must still be waiting (no slot available yet).
	select {
	case <-item2.waitCh:
		t.Fatal("item2 should not have been dispatched yet")
	default:
	}

	// Release the slot item1 holds so item2 gets dispatched.
	m.mu.Lock()
	m.releaseSlotLocked("g-fifo", "node-q1")
	m.mu.Unlock()

	select {
	case res := <-item2.waitCh:
		if res.err != nil {
			t.Fatalf("item2 expected dispatch, got error: %v", res.err)
		}
	case <-time.After(100 * time.Millisecond):
		t.Fatal("timeout: item2 was not dispatched after second slot release")
	}
}

// TestModelQueueOverflow verifies that admit rejects requests when the queue
// is at max_queue capacity.
func TestModelQueueOverflow(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-ov1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-ov1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill the node's capacity and set max_queue=1.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-overflow", groupPolicy{})
	g.policy.maxQueue = 1
	g.inflight["node-ov1"] = 1
	m.mu.Unlock()

	// Queue one item — should succeed.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	m.mu.Lock()
	m.groups["g-overflow"].queue = append(m.groups["g-overflow"].queue, item)
	m.mu.Unlock()

	// Second admit should fail immediately with errQueueFull.
	_, err := m.admit(context.Background(), "g-overflow", cands, groupPolicy{})
	if !errors.Is(err, errQueueFull) {
		t.Fatalf("expected errQueueFull, got: %v", err)
	}
}

// TestModelQueueTimeout verifies that queued items receive a timeout error
// when no slot becomes available before the deadline.
func TestModelQueueTimeout(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-to1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-to1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}

	m := newModelQueueManager(store)

	// Fill capacity and set a very short queue timeout.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-timeout", groupPolicy{})
	g.policy.queueTimeout = 20 * time.Millisecond
	g.inflight["node-to1"] = 1
	m.mu.Unlock()

	start := time.Now()
	_, err := m.admit(context.Background(), "g-timeout", cands, groupPolicy{})
	elapsed := time.Since(start)

	if !errors.Is(err, errQueueTimeout) {
		t.Fatalf("expected errQueueTimeout, got: %v", err)
	}
	if elapsed < 15*time.Millisecond {
		t.Fatalf("timed out too fast: %v", elapsed)
	}
}

// TestModelQueueTerminalReleaseDispatchesNext verifies that a terminal run event
// releases the in-flight slot and dispatches the next queued item.
func TestModelQueueTerminalReleaseDispatchesNext(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-tr1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-tr1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	bus := edgeevents.NewBus()
	m := newModelQueueManager(store)
	stop := m.startEventWatcher(bus)
	defer stop()

	// Fill capacity and record inflight.
	node, err := m.admit(context.Background(), "g-tr", cands, defPolicy)
	if err != nil {
		t.Fatalf("admit: %v", err)
	}
	m.trackInflight("g-tr", "run-tr-001", node.NodeID)

	// Queue a second item in a goroutine.
	resultCh := make(chan admitResult, 1)
	go func() {
		n, e := m.admit(context.Background(), "g-tr", cands, defPolicy)
		resultCh <- admitResult{node: n, err: e}
	}()

	waitForQueueLen(t, m, "g-tr", 1)

	// Terminal event for the first run — should unblock the queued admit.
	for _, termType := range []string{"complete", "error", "cancelled"} {
		t.Run(termType, func(t *testing.T) {
			// Reset state for each sub-test.
			m.mu.Lock()
			g := m.getOrCreateGroupLocked("g-tr-"+termType, groupPolicy{})
			g.inflight["node-tr1"] = 1
			item := &queueItem{
				candidates: cands,
				waitCh:     make(chan admitResult, 1),
				deadline:   time.Now().Add(2 * time.Second),
			}
			g.queue = []*queueItem{item}
			runID := "run-tr-" + termType
			m.inflightByRun[runID] = inflightRec{groupKey: "g-tr-" + termType, nodeID: "node-tr1"}
			m.mu.Unlock()

			bus.PublishRun(&iop.RunEvent{RunId: runID, Type: termType})

			select {
			case res := <-item.waitCh:
				if res.err != nil {
					t.Fatalf("expected dispatch, got error: %v", res.err)
				}
				if res.node == nil {
					t.Fatal("expected non-nil node")
				}
			case <-time.After(500 * time.Millisecond):
				t.Fatalf("timeout: queued item not dispatched after %q terminal event", termType)
			}
		})
	}

	// Cleanup the goroutine from the outer admit.
	m.mu.Lock()
	if g, ok := m.groups["g-tr"]; ok && len(g.queue) > 0 {
		g.inflight["node-tr1"] = 1
		m.tryDispatchLocked(g)
	}
	m.mu.Unlock()
	select {
	case <-resultCh:
	case <-time.After(100 * time.Millisecond):
	}
}

// TestModelQueueNodeDisconnectReleasesInflight verifies that a node disconnect
// event removes the node from queued item candidate lists, and that subsequent
// dispatch goes to a remaining live candidate (not the disconnected node).
func TestModelQueueNodeDisconnectReleasesInflight(t *testing.T) {
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd1",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})
	store.Add(&edgenode.NodeRecord{
		ID:      "node-nd2",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry1 := &edgenode.NodeEntry{NodeID: "node-nd1"}
	entry2 := &edgenode.NodeEntry{NodeID: "node-nd2"}
	// Queued item that can use either node (both at capacity 1).
	bothCands := []candidateNode{
		{entry: entry1, capacity: 1},
		{entry: entry2, capacity: 1},
	}

	bus := edgeevents.NewBus()
	m := newModelQueueManager(store)
	stop := m.startEventWatcher(bus)
	defer stop()

	// Fill node-nd1 and node-nd2 capacities and record inflight.
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-nd", groupPolicy{})
	g.inflight["node-nd1"] = 1
	g.inflight["node-nd2"] = 1
	m.inflightByRun["run-nd-x"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd1"}
	m.inflightByRun["run-nd-y"] = inflightRec{groupKey: "g-nd", nodeID: "node-nd2"}
	// Queue an item that can use either node.
	item := &queueItem{
		candidates: bothCands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(2 * time.Second),
	}
	g.queue = []*queueItem{item}
	m.mu.Unlock()

	// node-nd1 disconnects: its inflight is freed and it is removed from candidates.
	// node-nd2 is still full, so no dispatch yet.
	bus.PublishNode(&iop.EdgeNodeEvent{
		NodeId: "node-nd1",
		Type:   "node.disconnected",
	})

	// node-nd2's run terminates: now nd2 has capacity, dispatch goes to nd2.
	bus.PublishRun(&iop.RunEvent{RunId: "run-nd-y", Type: "complete"})

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after disconnect+terminal, got error: %v", res.err)
		}
		if res.node == nil {
			t.Fatal("expected non-nil node after disconnect release")
		}
		// Must be nd2 — nd1 was removed from candidates on disconnect.
		if res.node.NodeID != "node-nd2" {
			t.Errorf("expected dispatch to node-nd2, got %q (nd1 was disconnected)", res.node.NodeID)
		}
	case <-time.After(500 * time.Millisecond):
		t.Fatal("timeout: queued item not dispatched after node disconnect and terminal event")
	}

	// The inflight record for run-nd-x should be gone (removed at disconnect).
	m.mu.Lock()
	_, stillInFlight := m.inflightByRun["run-nd-x"]
	m.mu.Unlock()
	if stillInFlight {
		t.Error("run-nd-x should be removed from inflightByRun after node disconnect")
	}
}

// TestModelQueueUsesProviderCapacity verifies that the capacity supplied in the
// candidateNode (derived from adapter config) overrides the default of 1.
func TestModelQueueUsesProviderCapacity(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-pc1"}
	// Provider capacity = 2 (two concurrent slots on this node).
	cands := []candidateNode{{entry: entry, capacity: 2}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// First admit: inflight=0 < cap=2 → dispatched immediately.
	n1, err := m.admit(context.Background(), "g-pc", cands, defPolicy)
	if err != nil || n1 == nil {
		t.Fatalf("first admit: %v", err)
	}

	// Second admit: inflight=1 < cap=2 → still dispatched (not queued).
	n2, err := m.admit(context.Background(), "g-pc", cands, defPolicy)
	if err != nil || n2 == nil {
		t.Fatalf("second admit (capacity=2 should allow): %v", err)
	}

	// Third request must queue because inflight=2 == cap=2.
	item := &queueItem{
		candidates: cands,
		waitCh:     make(chan admitResult, 1),
		deadline:   time.Now().Add(200 * time.Millisecond),
	}
	m.mu.Lock()
	g := m.getOrCreateGroupLocked("g-pc", groupPolicy{})
	g.queue = append(g.queue, item)
	m.mu.Unlock()

	// Releasing one slot should dispatch the queued item.
	m.releaseSlot("g-pc", "node-pc1")

	select {
	case res := <-item.waitCh:
		if res.err != nil {
			t.Fatalf("expected dispatch after slot release, got: %v", res.err)
		}
		if res.node == nil || res.node.NodeID != "node-pc1" {
			t.Fatalf("unexpected node: %v", res.node)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: item not dispatched after slot release")
	}
}

// TestModelQueueUsesProviderQueuePolicy verifies that max_queue and
// queue_timeout from the policy parameter are respected.
func TestModelQueueUsesProviderQueuePolicy(t *testing.T) {
	t.Run("maxQueue enforced", func(t *testing.T) {
		entry := &edgenode.NodeEntry{NodeID: "node-pq1"}
		cands := []candidateNode{{entry: entry, capacity: 1}}
		// Provider policy: maxQueue=1.
		policy := groupPolicy{maxQueue: 1, queueTimeout: 5 * time.Second}

		m := newModelQueueManager(nil)

		// Fill capacity and queue one item manually.
		m.mu.Lock()
		g := m.getOrCreateGroupLocked("g-pq-max", policy)
		g.inflight["node-pq1"] = 1
		g.queue = []*queueItem{{
			candidates: cands,
			waitCh:     make(chan admitResult, 1),
			deadline:   time.Now().Add(5 * time.Second),
		}}
		m.mu.Unlock()

		// Second admit should fail: queue already at maxQueue=1.
		_, err := m.admit(context.Background(), "g-pq-max", cands, policy)
		if !errors.Is(err, errQueueFull) {
			t.Fatalf("expected errQueueFull, got: %v", err)
		}
	})

	t.Run("queueTimeout enforced", func(t *testing.T) {
		entry := &edgenode.NodeEntry{NodeID: "node-pq2"}
		cands := []candidateNode{{entry: entry, capacity: 1}}
		// Provider policy: very short timeout.
		policy := groupPolicy{maxQueue: 16, queueTimeout: 20 * time.Millisecond}

		m := newModelQueueManager(nil)

		// Fill capacity.
		m.mu.Lock()
		g := m.getOrCreateGroupLocked("g-pq-to", policy)
		g.inflight["node-pq2"] = 1
		m.mu.Unlock()

		start := time.Now()
		_, err := m.admit(context.Background(), "g-pq-to", cands, policy)
		elapsed := time.Since(start)

		if !errors.Is(err, errQueueTimeout) {
			t.Fatalf("expected errQueueTimeout, got: %v", err)
		}
		if elapsed < 15*time.Millisecond {
			t.Fatalf("timed out too fast: %v", elapsed)
		}
	})
}

// TestModelQueueContextCancelRemovesQueuedItem verifies that cancelling the
// context of a queued admit removes the item and returns context.Canceled.
func TestModelQueueContextCancelRemovesQueuedItem(t *testing.T) {
	entry := &edgenode.NodeEntry{NodeID: "node-cc1"}
	cands := []candidateNode{{entry: entry, capacity: 1}}
	defPolicy := groupPolicy{}

	m := newModelQueueManager(nil)

	// Fill the only slot.
	n, err := m.admit(context.Background(), "g-cc", cands, defPolicy)
	if err != nil || n == nil {
		t.Fatalf("initial admit: %v", err)
	}

	// Queue a second admit with a cancellable context.
	ctx, cancel := context.WithCancel(context.Background())
	resultCh := make(chan error, 1)
	go func() {
		_, err := m.admit(ctx, "g-cc", cands, defPolicy)
		resultCh <- err
	}()

	waitForQueueLen(t, m, "g-cc", 1)
	cancel()

	select {
	case err := <-resultCh:
		if !errors.Is(err, context.Canceled) {
			t.Errorf("expected context.Canceled, got: %v", err)
		}
	case <-time.After(200 * time.Millisecond):
		t.Fatal("timeout: context cancellation not handled")
	}

	// Queue must be empty after cancellation.
	m.mu.Lock()
	qLen := 0
	if g, ok := m.groups["g-cc"]; ok {
		qLen = len(g.queue)
	}
	m.mu.Unlock()
	if qLen != 0 {
		t.Errorf("queue should be empty after cancel, got %d items", qLen)
	}
}