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

package service

import (
	"context"
	"testing"

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

// reservationFixture admits one slot on a single-capacity provider and returns
// the manager, the reservation for that admission, and the slot key it holds.
func reservationFixture(t *testing.T, long bool, longCapacity int) (*modelQueueManager, *queueReservation, string) {
	t.Helper()
	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID:      "node-res",
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	entry := &edgenode.NodeEntry{NodeID: "node-res"}
	selected := &candidateNode{
		entry:               entry,
		capacity:            1,
		providerID:          "prov-1",
		longContextCapacity: longCapacity,
	}
	cands := []candidateNode{*selected}

	m := newModelQueueManager(store)
	admitted, _, err := m.admitWithReason(context.Background(), "g-res", "", "", cands, groupPolicy{}, long)
	if err != nil || admitted == nil {
		t.Fatalf("admit: %v", err)
	}

	return m, newQueueReservation(m, "g-res", selected, long), "node-res:prov-1"
}

func inflightCount(t *testing.T, m *modelQueueManager, groupKey, slot string) (int, int) {
	t.Helper()
	m.mu.Lock()
	defer m.mu.Unlock()
	group, ok := m.groups[groupKey]
	if !ok {
		t.Fatalf("group %q not found", groupKey)
	}
	return group.inflight[slot], group.longInflight[slot]
}

// TestQueueReservationReleasesExactlyOnce pins the release contract at each
// boundary a dispatch can fail on: the slot is freed once, and a repeated
// release cannot free a second slot.
func TestQueueReservationReleasesExactlyOnce(t *testing.T) {
	cases := []struct {
		name string
		// track reports whether the dispatch got as far as registering the run
		// inflight before failing.
		track  bool
		reason string
	}{
		{name: "prepare tunnel error before track", reason: "prepare-tunnel-error"},
		{name: "prepare run error before track", reason: "prepare-run-error"},
		{name: "build error before track", reason: "build-error"},
		{name: "unknown execution path before track", reason: "unknown-execution-path"},
		{name: "send error after track", track: true, reason: "send-error"},
		{name: "no event bus after track", track: true, reason: "no-event-bus"},
	}

	for _, tc := range cases {
		t.Run(tc.name, func(t *testing.T) {
			m, reservation, slot := reservationFixture(t, false, 0)

			if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 1 {
				t.Fatalf("expected 1 reserved slot after admission, got %d", inflight)
			}

			if tc.track {
				reservation.track("run-res")
			}

			reservation.release(tc.reason)
			if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
				t.Fatalf("expected the slot to be released, got %d in flight", inflight)
			}

			// A second release must not decrement a slot that now belongs to
			// another request.
			reservation.release(tc.reason)
			if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
				t.Fatalf("expected release to be idempotent, got %d in flight", inflight)
			}
			if tc.track {
				if _, ok := m.inflightByRun["run-res"]; ok {
					t.Fatal("expected the run to be removed from the inflight index")
				}
			}
		})
	}
}

// TestQueueReservationHandOffKeepsSlotHeld pins that a dispatched request keeps
// its slot: after hand-off the run/tunnel lifecycle owns the release, so a late
// release on the reservation must not free it.
func TestQueueReservationHandOffKeepsSlotHeld(t *testing.T) {
	m, reservation, slot := reservationFixture(t, false, 0)
	reservation.track("run-res")
	reservation.handOff()

	reservation.release("send-error")

	if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 1 {
		t.Fatalf("expected the dispatched run to keep its slot, got %d in flight", inflight)
	}

	// The lifecycle owner releases it exactly once, via the run registration.
	m.releaseRun("run-res", "complete")
	if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
		t.Fatalf("expected the terminal event to release the slot, got %d in flight", inflight)
	}
}

// TestQueueReservationLongSlotMatchesReservation pins that the long-context
// slot is released only when it was actually reserved: a long request on a
// provider that declares a long limit.
func TestQueueReservationLongSlotMatchesReservation(t *testing.T) {
	t.Run("long request on a long-capacity provider releases the long slot", func(t *testing.T) {
		m, reservation, slot := reservationFixture(t, true, 1)
		if _, longInflight := inflightCount(t, m, "g-res", slot); longInflight != 1 {
			t.Fatalf("expected 1 reserved long slot, got %d", longInflight)
		}

		reservation.release("send-error")

		inflight, longInflight := inflightCount(t, m, "g-res", slot)
		if inflight != 0 || longInflight != 0 {
			t.Fatalf("expected both slots released, got inflight=%d long=%d", inflight, longInflight)
		}
	})

	t.Run("long request without a declared long limit reserves no long slot", func(t *testing.T) {
		m, reservation, slot := reservationFixture(t, true, 0)
		if _, longInflight := inflightCount(t, m, "g-res", slot); longInflight != 0 {
			t.Fatalf("expected no long slot to be reserved, got %d", longInflight)
		}
		if reservation.long {
			t.Fatal("expected the reservation not to claim a long slot")
		}

		reservation.release("send-error")

		if inflight, _ := inflightCount(t, m, "g-res", slot); inflight != 0 {
			t.Fatalf("expected the slot to be released, got %d in flight", inflight)
		}
	})
}