[feat] Add Session-Aware Scheduling Logic

[SleepyLGod]

Pull Request Description

Current AIBrix Gateway uses a prefix-cache-aware router to route single LLM requests.

This PR is responsible for implementing session tracking and session-aware scheduling (and routing) within the Gateway.

Related Issues

Resolves: #633

Important: Before submitting, please complete the description above and review the checklist below.

---

Contribution Guidelines (Expand for Details)

We appreciate your contribution to aibrix! To ensure a smooth review process and maintain high code quality, please adhere to the following guidelines:

Pull Request Title Format

Your PR title should start with one of these prefixes to indicate the nature of the change:

• [Bug]: Corrections to existing functionality
• [CI]: Changes to build process or CI pipeline
• [Docs]: Updates or additions to documentation
• [API]: Modifications to aibrix's API or interface
• [CLI]: Changes or additions to the Command Line Interface
• [Misc]: For changes not covered above (use sparingly)

Note: For changes spanning multiple categories, use multiple prefixes in order of importance.

Submission Checklist

• PR title includes appropriate prefix(es)
• Changes are clearly explained in the PR description
• New and existing tests pass successfully
• Code adheres to project style and best practices
• Documentation updated to reflect changes (if applicable)
• Thorough testing completed, no regressions introduced

By submitting this PR, you confirm that you've read these guidelines and your changes align with the project's contribution standards.

[gemini-code-assist]

Summary of Changes

Hello @SleepyLGod, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request introduces core functionality for session tracking and session-aware scheduling within the AIBrix Gateway. It provides two distinct, thread-safe implementations for managing session state: one utilizing a sharded, channel-based approach for high concurrency, and another employing a traditional mutex-based design. This enables the Gateway to move beyond prefix-cache-aware routing for single requests to a more sophisticated session-aware scheduling system.

Highlights

• Session State Definition: Introduced a SessionState struct to encapsulate critical scheduling information for each session, including critical path service time, total wait time, pod affinity, and last activity timestamp.
• Sharded Concurrent Cache: Implemented ShardedSessionCache, a highly concurrent session cache that distributes session state across multiple goroutine-managed shards, communicating via channels to ensure thread safety and performance.
• Mutex-Based Cache: Provided MutexSessionCache as an alternative, simpler thread-safe session cache implementation using sync.RWMutex for protecting access to session data.
• Unified Cache API: Both cache implementations expose a consistent public API for retrieving session information (GetOrCreateForScheduler), updating session state (e.g., UpdateState for service time and wait time), and managing pod affinity (UpdateAffinity).
• Comprehensive Unit Tests: Added dedicated unit tests for both ShardedSessionCache and MutexSessionCache to validate their functionality, including single-threaded and concurrent update scenarios, ensuring correctness and robustness.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in issue comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments or fill out our survey to provide feedback.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request introduces two implementations for a session cache, one using sharding with channels (ShardedSessionCache) and another using a mutex (MutexSessionCache), to support session-aware scheduling. The implementations are well-structured and come with corresponding tests.

My review focuses on the ShardedSessionCache implementation, where I've identified several areas for improvement regarding concurrency safety, performance, and type safety. Specifically, the Close method has a race condition, any is used where specific types would be safer, there's an opportunity to optimize hashing with sync.Pool, and SessionState initialization could be more complete. I've also noted a minor correction for a comment in one of the test files. The MutexSessionCache implementation appears robust.

[Jeffwan]

Let's start reviewing this PR after v0.4.0 release

[Jeffwan]

@happyandslow can you help review the change first?