[feat] Add Session-Aware Scheduling Logic

pkg/plugins/gateway/scheduler/sessioninfo/cache.go

@@ -0,0 +1,87 @@
+/*
+Copyright 2025 The Aibrix Team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package sessioninfo
+
+import "time"
+
+// SessionCache is the interface for session state management.
+// It provides thread-safe operations for tracking session states
+// across multiple concurrent requests.
+//
+// Implementations:
+//   - MutexSessionCache: Simple mutex-based implementation for low to medium concurrency
+//   - ShardedSessionCache: High-performance sharded implementation for high concurrency
+type SessionCache interface {
+	// GetOrCreateForScheduler returns the inherited CST and total wait time
+	// for a new job in the given session. If the session doesn't exist,
+	// it will be created with zero values.
+	//
+	// This is the primary method used by the scheduler to get scheduling
+	// information for a new request.
+	//
+	// Returns:
+	//   - cst: The critical path service time of the session
+	//   - waitTime: The total accumulated wait time of the session
+	GetOrCreateForScheduler(sessionID string) (cst, waitTime time.Duration)
+
+	// UpdateState atomically updates the session state after a request completes.
+	//
+	// The update logic follows the ATLAS algorithm:
+	//   - TotalWaitTime is accumulated: totalWaitTime += waitTime
+	//   - CriticalPathServiceTime is updated to max(current, inheritedCST + executionTime)
+	//
+	// Parameters:
+	//   - sessionID: The session identifier
+	//   - inheritedCST: The CST value inherited when the request started
+	//   - executionTime: The actual execution time of this request
+	//   - waitTime: The time this request spent waiting in the queue
+	UpdateState(sessionID string, inheritedCST, executionTime, waitTime time.Duration)
+
+	// UpdateAffinity updates the pod affinity hint for a session.
+	// This can be used to optimize cache hits by routing subsequent
+	// requests from the same session to the same pod.
+	//
+	// This method also updates the LastActivityTimestamp to prevent
+	// the session from being cleaned up by the cleanup routine.
+	//
+	// Parameters:
+	//   - sessionID: The session identifier
+	//   - podName: The name of the pod to set as affinity hint
+	UpdateAffinity(sessionID, podName string)
+
+	// GetState retrieves a copy of the full session state.
+	// This method is primarily used for testing and debugging.
+	//
+	// Returns:
+	//   - state: A copy of the session state
+	//   - exists: false if the session doesn't exist
+	GetState(sessionID string) (state SessionState, exists bool)
+
+	// StartCleanupRoutine starts a background goroutine that periodically
+	// cleans up stale sessions that have been inactive for longer than timeout.
+	//
+	// The cleanup runs at the specified interval. Calling this method multiple
+	// times will start multiple cleanup routines (caller should avoid this).
+	//
+	// Parameters:
+	//   - interval: How often to run the cleanup
+	//   - timeout: Sessions inactive for longer than this will be removed
+	//
+	// Returns:
+	//   - stop: A function that stops the cleanup routine when called
+	StartCleanupRoutine(interval, timeout time.Duration) (stop func())
+}

pkg/plugins/gateway/scheduler/sessioninfo/cache_mutex.go

@@ -0,0 +1,158 @@
+/*
+Copyright 2025 The Aibrix Team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package sessioninfo
+
+import (
+	"sync"
+	"time"
+)
+
+// SessionState holds all the scheduling-relevant information for a single session
+type SessionState struct {
+	SessionID               string        // The session ID
+	CriticalPathServiceTime time.Duration // The critical path service time
+	TotalWaitTime           time.Duration // The total wait time (anti-starvation)
+	PodAffinity             string        // The pod affinity (later may needed)
+	LastActivityTimestamp   time.Time     // The last activity timestamp
+}
+
+// MutexSessionCache is a thread-safe, in-memory store for session states
+// using a sync.RWMutex.
+type MutexSessionCache struct {
+	mu       sync.RWMutex             // Protects the sessions map
+	sessions map[string]*SessionState // sessionID -> *SessionState
+}
+
+// NewMutexSessionCache creates a new in-memory session cache protected by a mutex.
+func NewMutexSessionCache() *MutexSessionCache {
+	return &MutexSessionCache{
+		sessions: make(map[string]*SessionState),
+	}
+}
+
+// getState is a private helper that assumes a write lock is already held.
+// It ensures a session state exists before any operation.
+func (sc *MutexSessionCache) getState(sessionID string) *SessionState {
+	state, exists := sc.sessions[sessionID]
+	if !exists {
+		state = &SessionState{
+			SessionID:             sessionID,
+			LastActivityTimestamp: time.Now(),
+		}
+		sc.sessions[sessionID] = state
+	}
+	return state
+}
+
+// GetState retrieves a copy of the state for a given sessionID
+// for read-only purposes.
+// It returns false if the session does not exist.
+func (sc *MutexSessionCache) GetState(sessionID string) (SessionState, bool) {
+	sc.mu.RLock()
+	defer sc.mu.RUnlock()
+
+	state, exists := sc.sessions[sessionID]
+	if !exists {
+		return SessionState{}, false
+	}
+
+	// Return a copy to ensure
+	// the caller cannot modify the internal state without a lock,
+	// which would cause a data race.
+	return *state, true
+}
+
+// GetOrCreateForScheduler is the primary method for the scheduler
+// to get the necessary info.
+// It returns the inherited CST and total wait time for a new job.
+func (sc *MutexSessionCache) GetOrCreateForScheduler(sessionID string) (
+	time.Duration, time.Duration) {
+	sc.mu.Lock() // Use a write lock because we might create a session.
+	defer sc.mu.Unlock()
+
+	state := sc.getState(sessionID)
+	return state.CriticalPathServiceTime, state.TotalWaitTime
+}
+
+// UpdateState atomically updates the session state after a request is finished.
+func (sc *MutexSessionCache) UpdateState(sessionID string, inheritedCST,
+	executionTime, waitTime time.Duration) {
+	sc.mu.Lock()
+	defer sc.mu.Unlock()
+
+	state := sc.getState(sessionID)
+
+	// Atomically update total wait time
+	state.TotalWaitTime += waitTime
+
+	// Atomically update CriticalPathServiceTime (ATLAS logic)
+	newPathLength := inheritedCST + executionTime
+	if newPathLength > state.CriticalPathServiceTime {
+		state.CriticalPathServiceTime = newPathLength
+	}
+
+	state.LastActivityTimestamp = time.Now()
+}
+
+// UpdateAffinity updates the pod affinity for a session.
+func (sc *MutexSessionCache) UpdateAffinity(sessionID, podName string) {
+	sc.mu.Lock()
+	defer sc.mu.Unlock()
+
+	state := sc.getState(sessionID)
+	state.PodAffinity = podName
+	state.LastActivityTimestamp = time.Now()
+}
+
+// StartCleanupRoutine starts a background goroutine that periodically
+// cleans up stale sessions.
+// It returns a function that can be called to stop the routine.
+func (sc *MutexSessionCache) StartCleanupRoutine(interval,
+	timeout time.Duration) (stop func()) {
+	ticker := time.NewTicker(interval)
+	done := make(chan struct{})
+
+	go func() {
+		for {
+			select {
+			case <-ticker.C:
+				sc.cleanup(timeout)
+			case <-done:
+				ticker.Stop()
+				return
+			}
+		}
+	}()
+
+	return func() {
+		close(done)
+	}
+}
+
+// cleanup removes sessions that have been inactive for longer than the timeout.
+// This is a private method that assumes the caller handles locking.
+func (sc *MutexSessionCache) cleanup(timeout time.Duration) {
+	sc.mu.Lock()
+	defer sc.mu.Unlock()
+
+	now := time.Now()
+	for sessionID, state := range sc.sessions {
+		if now.Sub(state.LastActivityTimestamp) > timeout {
+			delete(sc.sessions, sessionID)
+		}
+	}
+}

pkg/plugins/gateway/scheduler/sessioninfo/cache_mutex_test.go

@@ -0,0 +1,134 @@
+/*
+Copyright 2025 The Aibrix Team.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+	http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+*/
+
+package sessioninfo
+
+import (
+	"fmt"
+	"sync"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+)
+
+// TestMutexCache_GetOrCreateForScheduler_NewSession tests the GetOrCreateForScheduler method.
+func TestMutexCache_GetOrCreateForScheduler_NewSession(t *testing.T) {
+	cache := NewMutexSessionCache()
+	cst, waitTime := cache.GetOrCreateForScheduler("session1")
+
+	assert.Equal(t, time.Duration(0), cst)
+	assert.Equal(t, time.Duration(0), waitTime)
+}
+
+// TestMutexCache_UpdateState_Single tests the UpdateState method.
+func TestMutexCache_UpdateState_Single(t *testing.T) {
+	cache := NewMutexSessionCache()
+
+	// First update (like a serial request)
+	cache.UpdateState("session1", 0, 5*time.Second, 2*time.Second)
+	state, _ := cache.GetState("session1")
+	assert.Equal(t, 5*time.Second, state.CriticalPathServiceTime)
+	assert.Equal(t, 2*time.Second, state.TotalWaitTime)
+
+	// Second update (another serial request)
+	// InheritedCST should be the CST from the previous state (5s)
+	cache.UpdateState("session1", 5*time.Second, 3*time.Second, 1*time.Second)
+	state, _ = cache.GetState("session1")
+	assert.Equal(t, 8*time.Second, state.CriticalPathServiceTime) // 5s + 3s
+	assert.Equal(t, 3*time.Second, state.TotalWaitTime)           // 2s + 1s
+}
+
+// TestMutexCache_UpdateState_Concurrent tests the UpdateState method.
+func TestMutexCache_UpdateState_Concurrent(t *testing.T) {
+	cache := NewMutexSessionCache()
+	concurrency := 1000
+	var wg sync.WaitGroup
+	wg.Add(concurrency)
+
+	// Simulate 100 parallel requests for the same session finishing.
+	// All inherited CST=0, as they started when the session's CST was 0.
+	for i := 0; i < concurrency; i++ {
+		go func(execTimeMs int) {
+			defer wg.Done()
+			cache.UpdateState("session1", 0,
+				time.Duration(execTimeMs)*time.Millisecond,
+				10*time.Millisecond)
+		}(i + 1)
+	}
+	wg.Wait()
+
+	state, exists := cache.GetState("session1")
+	assert.True(t, exists)
+
+	// Final CST should be the max of all new path lengths,
+	// which is max(0+1ms, 0+2ms, ... 0+100ms, ..., 0+1000ms) = 1000ms
+	assert.Equal(t, 1000*time.Millisecond, state.CriticalPathServiceTime)
+
+	// Total wait time should be the sum of all wait times:
+	// 1000 * 10ms = 10000ms
+	assert.Equal(t, 10000*time.Millisecond, state.TotalWaitTime)
+}
+
+// TestMutexCache_UpdateAffinity_Concurrent tests the UpdateAffinity method.
+func TestMutexCache_UpdateAffinity_Concurrent(t *testing.T) {
+	cache := NewMutexSessionCache()
+	concurrency := 10
+	var wg sync.WaitGroup
+	wg.Add(concurrency)
+
+	for i := 0; i < concurrency; i++ {
+		go func(podNum int) {
+			defer wg.Done()
+			cache.UpdateAffinity("session1",
+				fmt.Sprintf("pod%d", podNum))
+		}(i)
+	}
+	wg.Wait()
+
+	state, exists := cache.GetState("session1")
+	assert.True(t, exists)
+	// Due to the race, we can't know the final value,
+	// but it must be one of the values we set.
+	assert.Contains(t, []string{"pod0", "pod1", "pod2", "pod3",
+		"pod4", "pod5", "pod6", "pod7", "pod8", "pod9"},
+		state.PodAffinity)
+}
+
+// TestMutexCache_Cleanup tests the Cleanup method.
+func TestMutexCache_Cleanup(t *testing.T) {
+	cache := NewMutexSessionCache()
+
+	// Create session1
+	cache.UpdateState("session1", 0, 1*time.Second, 0)
+
+	// Wait for 2 seconds, making session1 stale relative to a 1.5s timeout
+	time.Sleep(2 * time.Second)
+
+	// Create/update session2, making it fresh
+	cache.UpdateState("session2", 0, 1*time.Second, 0)
+
+	// Now, cleanup sessions older than 1.5 seconds
+	cache.cleanup(1500 * time.Millisecond)
+
+	// session1 should be gone because it's ~2 seconds old
+	_, exists := cache.GetState("session1")
+	assert.False(t, exists, "session1 should be stale and cleaned up")
+
+	// session2 should still exist because it's very fresh
+	_, exists = cache.GetState("session2")
+	assert.True(t, exists, "session2 should be fresh and remain")
+}