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

package service

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

	toki "git.toki-labs.com/toki/proto-socket/go"
	"google.golang.org/protobuf/proto"

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

func TestSubmitRunModelQueueUsesRoutePolicyBeforeProviderInstancePolicy(t *testing.T) {
	parserMap := toki.ParserMap{
		toki.TypeNameOf(&iop.RunRequest{}): func(b []byte) (proto.Message, error) {
			m := &iop.RunRequest{}
			return m, proto.Unmarshal(b, m)
		},
	}

	edgeConn, nodeConn := net.Pipe()
	defer edgeConn.Close()
	defer nodeConn.Close()

	edgeClient := toki.NewTcpClient(edgeConn, 0, 0, parserMap)
	nodeClient := toki.NewTcpClient(nodeConn, 0, 0, parserMap)
	toki.AddListenerTyped[*iop.RunRequest](&nodeClient.Communicator, func(*iop.RunRequest) {})

	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "prov-node-route-policy", Client: edgeClient})

	store := edgenode.NewNodeStore()
	store.Add(&edgenode.NodeRecord{
		ID: "prov-node-route-policy",
		Adapters: config.AdaptersConf{
			OllamaInstances: []config.OllamaInstanceConf{
				{
					Name:     "ollama-local",
					Enabled:  true,
					Capacity: 1,
					MaxQueue: 10,
				},
			},
		},
		Runtime: config.RuntimeConf{Concurrency: 1},
	})

	bus := edgeevents.NewBus()
	svc := New(reg, bus)
	svc.SetNodeStore(store)

	// run 1: fills capacity (capacity is 1)
	res1, err := svc.SubmitRun(context.Background(), SubmitRunRequest{
		ModelGroupKey: "route-policy-group",
		Adapter:       "ollama-local",
		Background:    true,
		MaxQueue:      1,
	})
	if err != nil {
		t.Fatalf("run1 error: %v", err)
	}
	defer res1.Close()

	// run 2: enters queue in a separate goroutine
	ctx2, cancel2 := context.WithCancel(context.Background())
	defer cancel2()
	errCh2 := make(chan error, 1)
	go func() {
		res2, err := svc.SubmitRun(ctx2, SubmitRunRequest{
			ModelGroupKey: "route-policy-group",
			Adapter:       "ollama-local",
			Background:    true,
			MaxQueue:      1,
		})
		if err == nil {
			res2.Close()
		}
		errCh2 <- err
	}()

	// Wait explicitly until run 2 enters the queue
	waitForQueueLen(t, svc.queue, "route-policy-group", 1)

	// run 3: should be rejected immediately because queue is full (max queue = 1, and run2 is in the queue)
	// We use a short timeout context to prevent blocking indefinitely if something fails.
	ctx3, cancel3 := context.WithTimeout(context.Background(), 2*time.Second)
	defer cancel3()

	_, err3 := svc.SubmitRun(ctx3, SubmitRunRequest{
		ModelGroupKey: "route-policy-group",
		Adapter:       "ollama-local",
		Background:    true,
		MaxQueue:      1,
	})
	if err3 == nil {
		t.Fatal("expected third run to fail immediately due to queue full")
	}
	if !strings.Contains(err3.Error(), "queue is full") {
		t.Errorf("expected queue is full error, got: %v", err3)
	}

	// clean up run2 and verify it exits
	cancel2()
	select {
	case err2 := <-errCh2:
		if err2 == nil {
			t.Error("expected run2 to fail with context cancelled, got nil")
		} else if !errors.Is(err2, context.Canceled) {
			t.Errorf("expected run2 to fail with context cancelled, got: %v", err2)
		}
	case <-time.After(1 * time.Second):
		t.Error("timeout waiting for run2 goroutine to exit")
	}
}

// TestRefreshProviderCapacityAffectsNextDispatch verifies that a provider
// capacity change applied via SetRuntimeConfig is reflected in the candidate
// capacity used for the next dispatch admission, since candidates are rebuilt
// from the live runtime snapshot on each request.
func TestRefreshProviderCapacityAffectsNextDispatch(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-cap"})
	svc := New(reg, nil)

	newStore := func(capacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID: "node-cap",
			Providers: []config.NodeProviderConf{
				{
					ID:       "prov-a",
					Type:     "vllm",
					Category: config.CategoryAPI,
					Adapter:  "vllm-gpu",
					Models:   []string{"served-a"},
					Health:   "available",
					Capacity: capacity,
				},
			},
		})
		return store
	}
	catalog := []config.ModelCatalogEntry{{
		ID:        "qwen3.6:35b",
		Providers: map[string]string{"prov-a": "served-a"},
	}}

	req := SubmitRunRequest{ModelGroupKey: "qwen3.6:35b", ProviderPool: true}

	svc.SetRuntimeConfig(newStore(2), catalog)
	store, cat := svc.runtimeConfigSnapshot()
	cands, _, err := svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve before refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].capacity != 2 {
		t.Fatalf("before refresh: got candidates=%+v, want one with capacity=2", cands)
	}

	// Refresh raises capacity to 8; the next candidate build must observe it.
	svc.SetRuntimeConfig(newStore(8), catalog)
	store, cat = svc.runtimeConfigSnapshot()
	cands, _, err = svc.resolveProviderPoolCandidates(req, store, cat)
	if err != nil {
		t.Fatalf("resolve after refresh: %v", err)
	}
	if len(cands) != 1 || cands[0].capacity != 8 {
		t.Fatalf("after refresh: got candidates=%+v, want one with capacity=8", cands)
	}
}

func TestRuntimeConfigSnapshotConcurrentReplace(t *testing.T) {
	reg := edgenode.NewRegistry()
	reg.Register(&edgenode.NodeEntry{NodeID: "node-refresh"})
	svc := New(reg, nil)

	newStore := func(providerID, servedModel string, capacity int) *edgenode.NodeStore {
		store := edgenode.NewNodeStore()
		store.Add(&edgenode.NodeRecord{
			ID:      "node-refresh",
			Runtime: config.RuntimeConf{Concurrency: capacity},
			Providers: []config.NodeProviderConf{
				{
					ID:       providerID,
					Type:     "vllm",
					Category: config.CategoryAPI,
					Adapter:  "vllm-gpu",
					Models:   []string{servedModel},
					Health:   "available",
					Capacity: capacity,
				},
			},
		})
		return store
	}
	newCatalog := func(providerID, servedModel string) []config.ModelCatalogEntry {
		return []config.ModelCatalogEntry{
			{
				ID: "qwen3.6:35b",
				Providers: map[string]string{
					providerID: servedModel,
				},
			},
		}
	}

	storeA := newStore("prov-a", "served-a", 2)
	catalogA := newCatalog("prov-a", "served-a")
	storeB := newStore("prov-b", "served-b", 4)
	catalogB := newCatalog("prov-b", "served-b")
	svc.SetRuntimeConfig(storeA, catalogA)

	const readers = 8
	const iterations = 200
	start := make(chan struct{})
	errCh := make(chan error, readers)
	var wg sync.WaitGroup

	for i := 0; i < readers; i++ {
		wg.Add(1)
		go func() {
			defer wg.Done()
			<-start
			for j := 0; j < iterations; j++ {
				store, catalog := svc.runtimeConfigSnapshot()
				candidates, _, err := svc.resolveProviderPoolCandidates(SubmitRunRequest{
					ModelGroupKey: "qwen3.6:35b",
					ProviderPool:  true,
				}, store, catalog)
				if err != nil {
					errCh <- err
					return
				}
				if len(candidates) != 1 {
					errCh <- errors.New("expected exactly one provider candidate")
					return
				}
				_ = svc.ListNodeSnapshots()
			}
		}()
	}

	wg.Add(1)
	go func() {
		defer wg.Done()
		<-start
		for i := 0; i < iterations; i++ {
			svc.SetRuntimeConfig(storeA, catalogA)
			svc.SetRuntimeConfig(storeB, catalogB)
		}
	}()

	close(start)
	wg.Wait()
	close(errCh)
	for err := range errCh {
		if err != nil {
			t.Fatalf("runtime config snapshot reader observed inconsistent state: %v", err)
		}
	}
}