fix: apply scheduling strategies by TID with TGID fallback

plugin/gthulhu/gthulhu.go

@@ -66,12 +66,12 @@ type GthulhuPlugin struct {
 	// Global vruntime
 	minVruntime uint64
 
-	// Strategy map for PID-based scheduling strategies
-	oldStrategyMap  map[int32]util.SchedulingStrategy
-	strategyMap     map[int32]util.SchedulingStrategy
-	newStrategy     []util.SchedulingStrategy
-	removedStrategy []util.SchedulingStrategy
-	strategyMu      sync.RWMutex
+	// strategyMap is the latest desired strategy set (keyed by task id);
+	// appliedStrategyMap is the set last handed to the scheduler, so
+	// GetChangedStrategies can diff the two into a coalesced changed/removed set.
+	strategyMap        map[int32]util.SchedulingStrategy
+	appliedStrategyMap map[int32]util.SchedulingStrategy
+	strategyMu         sync.RWMutex
 
 	// JWT client for API authentication
 	jwtClient *JWTClient
@@ -82,12 +82,13 @@ type GthulhuPlugin struct {
 
 func NewGthulhuPlugin(sliceNsDefault, sliceNsMin uint64) *GthulhuPlugin {
 	plugin := &GthulhuPlugin{
-		sliceNsDefault: 5000 * 1000, // 5ms (default)
-		sliceNsMin:     500 * 1000,  // 0.5ms (default)
-		taskPool:       make([]Task, taskPoolSize),
-		taskPoolCount:  0,
-		minVruntime:    0,
-		strategyMap:    make(map[int32]util.SchedulingStrategy),
+		sliceNsDefault:     5000 * 1000, // 5ms (default)
+		sliceNsMin:         500 * 1000,  // 0.5ms (default)
+		taskPool:           make([]Task, taskPoolSize),
+		taskPoolCount:      0,
+		minVruntime:        0,
+		strategyMap:        make(map[int32]util.SchedulingStrategy),
+		appliedStrategyMap: make(map[int32]util.SchedulingStrategy),
 	}
 
 	// Override defaults if provided
@@ -130,7 +131,7 @@ func (g *GthulhuPlugin) SelectCPU(s reg.Sched, t *models.QueuedTask) (error, int
 }
 
 func (g *GthulhuPlugin) DetermineTimeSlice(s reg.Sched, t *models.QueuedTask) uint64 {
-	return g.getTaskExecutionTime(t.Pid)
+	return g.getTaskExecutionTime(t)
 }
 
 func (g *GthulhuPlugin) GetPoolCount() uint64 {
@@ -281,33 +282,49 @@ func lessQueuedTask(a, b *Task) bool {
 	return a.QueuedTask.Pid < b.QueuedTask.Pid
 }
 
-// applySchedulingStrategy applies scheduling strategies to a task
+// applySchedulingStrategy gives a task minimum vtime when a matching strategy
+// boosts it (Priority > 0), and reports whether it was boosted. A strategy that
+// only sets a custom time slice (Priority == 0) must not jump the run queue, so
+// it reports false here; getTaskExecutionTime still supplies its slice.
 func (g *GthulhuPlugin) applySchedulingStrategy(task *models.QueuedTask) bool {
-	g.strategyMu.RLock()
-	strategy, exists := g.strategyMap[task.Tgid]
-	g.strategyMu.RUnlock()
-	if exists {
-		// Apply strategy
-		if strategy.Priority > 0 {
-			// Priority tasks get minimum vtime
-			task.Vtime = 0
-		}
-		return true
+	strategy, exists := g.lookupTaskStrategy(task)
+	if !exists || strategy.Priority <= 0 {
+		return false
 	}
-	return false
+	task.Vtime = 0
+	return true
 }
 
-// getTaskExecutionTime returns the custom execution time for a task if defined
-func (g *GthulhuPlugin) getTaskExecutionTime(pid int32) uint64 {
-	g.strategyMu.RLock()
-	strategy, exists := g.strategyMap[pid]
-	g.strategyMu.RUnlock()
+// getTaskExecutionTime returns the custom time slice for a task, or 0 when no
+// matching strategy defines one.
+func (g *GthulhuPlugin) getTaskExecutionTime(task *models.QueuedTask) uint64 {
+	strategy, exists := g.lookupTaskStrategy(task)
 	if exists && strategy.ExecutionTime > 0 {
 		return strategy.ExecutionTime
 	}
 	return 0
 }
 
+// lookupTaskStrategy returns the strategy for a task, preferring an exact
+// thread (TID) match over a thread-group (TGID) match. Node policies key by
+// TID, so a thread-specific rule wins; Pod policies key by the group leader's
+// PID, so every thread of the group still resolves through the TGID fallback.
+// Priority and time-slice must share this lookup, or one strategy would act
+// differently between the two paths.
+func (g *GthulhuPlugin) lookupTaskStrategy(task *models.QueuedTask) (util.SchedulingStrategy, bool) {
+	g.strategyMu.RLock()
+	defer g.strategyMu.RUnlock()
+	if strategy, ok := g.strategyMap[task.Pid]; ok {
+		return strategy, true
+	}
+	if task.Tgid != task.Pid {
+		if strategy, ok := g.strategyMap[task.Tgid]; ok {
+			return strategy, true
+		}
+	}
+	return util.SchedulingStrategy{}, false
+}
+
 // InitJWTClient initializes the JWT client for API authentication
 func (g *GthulhuPlugin) InitJWTClient(
 	publicKeyPath,
@@ -365,59 +382,46 @@ func (g *GthulhuPlugin) FetchSchedulingStrategies(apiUrl string) ([]util.Schedul
 	return fetchSchedulingStrategies(g.jwtClient, apiUrl)
 }
 
-// UpdateStrategyMap updates the strategy map from a slice of strategies
+// UpdateStrategyMap replaces the desired strategy set. The changed/removed diff
+// is computed later in GetChangedStrategies against the last applied set, so
+// intermediate churn (e.g. a strategy removed then re-added before the next
+// drain) coalesces to the correct final state instead of a stale event stream.
 func (g *GthulhuPlugin) UpdateStrategyMap(strategies []util.SchedulingStrategy) {
-	// Create a new map to avoid concurrent access issues
 	newMap := make(map[int32]util.SchedulingStrategy)
-
 	for _, strategy := range strategies {
 		newMap[int32(strategy.PID)] = strategy
 	}
-
-	// Replace the old map with the new one
 	g.strategyMu.Lock()
-	g.oldStrategyMap = g.strategyMap
 	g.strategyMap = newMap
-	changed, removed := g.caculateChangedStrategies()
-	g.newStrategy = append(g.newStrategy, changed...)
-	g.removedStrategy = append(g.removedStrategy, removed...)
 	g.strategyMu.Unlock()
 }
 
-// Campare g.oldStrategyMap and g.strategyMap and return the list of SchedulingStrategy that have changed strategies
-func (g *GthulhuPlugin) caculateChangedStrategies() ([]util.SchedulingStrategy, []util.SchedulingStrategy) {
+// GetChangedStrategies returns the strategies to apply (changed or new) and to
+// remove so the scheduler's applied set matches the current desired set, then
+// records the current set as applied. Diffing against the last applied set (not
+// a running event queue) guarantees a strategy is never in both lists, so a
+// remove-then-re-add between drains is not mistaken for a deletion.
+func (g *GthulhuPlugin) GetChangedStrategies() ([]util.SchedulingStrategy, []util.SchedulingStrategy) {
 	changed := []util.SchedulingStrategy{}
 	removed := []util.SchedulingStrategy{}
 
-	// Check for removed strategies
-	for pid, oldStrategy := range g.oldStrategyMap {
-		_, exists := g.strategyMap[pid]
-		if !exists {
-			removed = append(removed, oldStrategy)
+	g.strategyMu.Lock()
+	defer g.strategyMu.Unlock()
+
+	for pid, strategy := range g.strategyMap {
+		if applied, ok := g.appliedStrategyMap[pid]; !ok || applied != strategy {
+			changed = append(changed, strategy)
 		}
 	}
-
-	// Check for changed or new strategies
-	for pid, newStrategy := range g.strategyMap {
-		oldStrategy, exists := g.oldStrategyMap[pid]
-		if !exists || oldStrategy != newStrategy {
-			changed = append(changed, newStrategy)
+	for pid, applied := range g.appliedStrategyMap {
+		if _, ok := g.strategyMap[pid]; !ok {
+			removed = append(removed, applied)
 		}
 	}
-	return changed, removed
-}
 
-// Campare g.oldStrategyMap and g.strategyMap and return the list of SchedulingStrategy that have changed strategies
-func (g *GthulhuPlugin) GetChangedStrategies() ([]util.SchedulingStrategy, []util.SchedulingStrategy) {
-	changed := []util.SchedulingStrategy{}
-	removed := []util.SchedulingStrategy{}
-	g.strategyMu.RLock()
-	defer g.strategyMu.RUnlock()
-
-	// copy g.newStrategy to changed and clear g.newStrategy
-	changed = append(changed, g.newStrategy...)
-	removed = append(removed, g.removedStrategy...)
-	g.newStrategy = []util.SchedulingStrategy{}
-	g.removedStrategy = []util.SchedulingStrategy{}
+	g.appliedStrategyMap = make(map[int32]util.SchedulingStrategy, len(g.strategyMap))
+	for pid, strategy := range g.strategyMap {
+		g.appliedStrategyMap[pid] = strategy
+	}
 	return changed, removed
 }

plugin/gthulhu/gthulhu_test.go

@@ -1,13 +1,19 @@
 package gthulhu
 
 import (
+	"sync"
 	"testing"
 
 	"github.com/Gthulhu/plugin/models"
 	reg "github.com/Gthulhu/plugin/plugin/internal/registry"
 	"github.com/Gthulhu/plugin/plugin/util"
 )
 
+// makeTask builds a QueuedTask identified by its thread (TID) and group (TGID).
+func makeTask(tid, tgid int32) *models.QueuedTask {
+	return &models.QueuedTask{Pid: tid, Tgid: tgid}
+}
+
 // TestGthulhuPluginInstanceIsolation verifies that multiple GthulhuPlugin instances maintain independent state
 func TestGthulhuPluginInstanceIsolation(t *testing.T) {
 	// Create two instances with different configurations
@@ -142,6 +148,163 @@ func TestGthulhuPluginUpdateStrategyMap(t *testing.T) {
 	}
 }
 
+// TestStrategyLookupPrefersTID verifies a TID-keyed strategy binds to the exact
+// thread for both priority and time slice and does not leak to a sibling thread.
+func TestStrategyLookupPrefersTID(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	// Node policy on a single worker thread (tid 501) of process 500.
+	g.UpdateStrategyMap([]util.SchedulingStrategy{
+		{PID: 501, Priority: 1, ExecutionTime: 7000},
+	})
+
+	worker := makeTask(501, 500)
+	if !g.applySchedulingStrategy(worker) {
+		t.Fatal("worker thread should match its TID-keyed strategy")
+	}
+	if worker.Vtime != 0 {
+		t.Errorf("priority worker Vtime = %d; want 0", worker.Vtime)
+	}
+	if got := g.getTaskExecutionTime(worker); got != 7000 {
+		t.Errorf("worker execution time = %d; want 7000", got)
+	}
+
+	// A sibling thread of the same group must not inherit a TID-specific rule.
+	sibling := makeTask(502, 500)
+	if g.applySchedulingStrategy(sibling) {
+		t.Error("sibling thread must not match a TID-specific strategy")
+	}
+	if got := g.getTaskExecutionTime(sibling); got != 0 {
+		t.Errorf("sibling execution time = %d; want 0", got)
+	}
+}
+
+// TestStrategyLookupTGIDFallback verifies a group-leader-keyed strategy (a Pod
+// policy) reaches every thread of the group via the TGID fallback, for both
+// priority and time slice.
+func TestStrategyLookupTGIDFallback(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{
+		{PID: 800, Priority: 1, ExecutionTime: 9000},
+	})
+
+	for _, tid := range []int32{800, 801, 802} {
+		task := makeTask(tid, 800)
+		if !g.applySchedulingStrategy(task) {
+			t.Fatalf("thread %d should match via TGID fallback", tid)
+		}
+		if got := g.getTaskExecutionTime(task); got != 9000 {
+			t.Errorf("thread %d execution time = %d; want 9000", tid, got)
+		}
+	}
+}
+
+// TestStrategyLookupTIDWinsOverTGID verifies a thread-specific strategy takes
+// precedence over a group-wide one for that thread.
+func TestStrategyLookupTIDWinsOverTGID(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{
+		{PID: 900, Priority: 1, ExecutionTime: 1000}, // whole group
+		{PID: 901, Priority: 1, ExecutionTime: 2000}, // thread 901 only
+	})
+	if got := g.getTaskExecutionTime(makeTask(901, 900)); got != 2000 {
+		t.Errorf("thread 901 execution time = %d; want 2000 (TID wins)", got)
+	}
+	if got := g.getTaskExecutionTime(makeTask(902, 900)); got != 1000 {
+		t.Errorf("thread 902 execution time = %d; want 1000 (TGID fallback)", got)
+	}
+}
+
+// TestPodPolicySliceAppliesToAllGroupThreads pins a deliberate behavior: a Pod
+// policy keyed by the group leader PID applies its custom time slice to every
+// thread of the group, not just the leader, so the slice fans out the same way
+// priority already does via the TGID. This is an intentional contract - do not
+// narrow it back to leader-only.
+func TestPodPolicySliceAppliesToAllGroupThreads(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	// Pod policy on leader PID 700, priority left off to isolate the slice.
+	g.UpdateStrategyMap([]util.SchedulingStrategy{
+		{PID: 700, Priority: 0, ExecutionTime: 12345},
+	})
+	if got := g.getTaskExecutionTime(makeTask(700, 700)); got != 12345 {
+		t.Errorf("leader slice = %d; want 12345", got)
+	}
+	if got := g.getTaskExecutionTime(makeTask(701, 700)); got != 12345 {
+		t.Errorf("non-leader thread slice = %d; want 12345 (fans to whole group)", got)
+	}
+}
+
+// TestSliceOnlyStrategyDoesNotJumpQueue verifies a Priority==0 strategy supplies
+// its custom time slice but is not boosted to the front of the run queue.
+func TestSliceOnlyStrategyDoesNotJumpQueue(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{
+		{PID: 600, Priority: 0, ExecutionTime: 4000},
+	})
+	task := makeTask(600, 600)
+	task.Vtime = 999 // sentinel: a non-boost must not reset vtime to 0
+	if g.applySchedulingStrategy(task) {
+		t.Error("a Priority==0 strategy must not report a boost (would force Deadline 0)")
+	}
+	if task.Vtime != 999 {
+		t.Errorf("non-boost strategy changed Vtime to %d; want 999 untouched", task.Vtime)
+	}
+	if got := g.getTaskExecutionTime(task); got != 4000 {
+		t.Errorf("slice-only strategy execution time = %d; want 4000", got)
+	}
+}
+
+// TestGetChangedStrategiesCoalescesRemoveThenReadd verifies that removing a PID
+// and re-adding it (with a new value) before a drain yields one changed entry
+// and no removal - the scheduler must not end up deleting the re-added key.
+func TestGetChangedStrategiesCoalescesRemoveThenReadd(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{{PID: 42, Priority: 1, ExecutionTime: 100}})
+	if changed, removed := g.GetChangedStrategies(); len(changed) != 1 || len(removed) != 0 {
+		t.Fatalf("baseline drain: changed=%d removed=%d; want 1,0", len(changed), len(removed))
+	}
+	// Remove 42, then re-add it with a new value, both before the next drain.
+	g.UpdateStrategyMap(nil)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{{PID: 42, Priority: 2, ExecutionTime: 200}})
+
+	changed, removed := g.GetChangedStrategies()
+	if len(removed) != 0 {
+		t.Errorf("removed=%v; want empty (42 must not be deleted)", removed)
+	}
+	if len(changed) != 1 || changed[0].PID != 42 || changed[0].ExecutionTime != 200 {
+		t.Errorf("changed=%v; want a single 42 with ExecutionTime 200", changed)
+	}
+}
+
+// TestGetChangedStrategiesCoalescesAddThenRemove verifies a PID added then
+// removed before a drain nets to nothing.
+func TestGetChangedStrategiesCoalescesAddThenRemove(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{{PID: 7, Priority: 1}})
+	g.UpdateStrategyMap(nil)
+	if changed, removed := g.GetChangedStrategies(); len(changed) != 0 || len(removed) != 0 {
+		t.Errorf("changed=%d removed=%d; want 0,0 (add then remove nets to nothing)", len(changed), len(removed))
+	}
+}
+
+// TestGetChangedStrategiesConcurrent drains the change queues from several
+// goroutines at once so the race detector guards the write-lock fix.
+func TestGetChangedStrategiesConcurrent(t *testing.T) {
+	g := NewGthulhuPlugin(0, 0)
+	g.UpdateStrategyMap([]util.SchedulingStrategy{{PID: 1, Priority: 1}})
+
+	var wg sync.WaitGroup
+	for i := 0; i < 4; i++ {
+		wg.Add(1)
+		go func() {
+			defer wg.Done()
+			for j := 0; j < 500; j++ {
+				g.GetChangedStrategies()
+			}
+		}()
+	}
+	wg.Wait()
+}
+
 // MockScheduler implements the plugin.Sched interface for testing
 type MockScheduler struct {
 	taskQueue     []*models.QueuedTask