Integrate session manager in router

Signed-off-by: VanderChen <vanderchen@outlook.com>

Author: VanderChen

cmd/router/main.go

@@ -52,13 +52,10 @@ func main() {
 		log.Fatalf("Failed to create Router API server: %v", err)
 	}
 
-	// Setup signal handling
-	ctx, cancel := context.WithCancel(context.Background())
+	// Setup signal handling with context cancellation
+	ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt, syscall.SIGTERM)
 	defer cancel()
 
-	sigCh := make(chan os.Signal, 1)
-	signal.Notify(sigCh, os.Interrupt, syscall.SIGTERM)
-
 	// Start Router API server in goroutine
 	errCh := make(chan error, 1)
 	go func() {
@@ -70,10 +67,10 @@ func main() {
 
 	// Wait for signal or error
 	select {
-	case <-sigCh:
+	case <-ctx.Done():
 		log.Println("Received shutdown signal, shutting down gracefully...")
-		cancel()
-		time.Sleep(2 * time.Second) // Give server time to shutdown gracefully
+		// Wait for server to finish shutting down
+		time.Sleep(2 * time.Second)
 	case err := <-errCh:
 		log.Fatalf("Server error: %v", err)
 	}

pkg/router/apiserver.go

@@ -5,9 +5,12 @@ import (
 	"fmt"
 	"log"
 	"net/http"
+	"os"
 	"time"
 
 	"github.com/gin-gonic/gin"
+	redisv9 "github.com/redis/go-redis/v9"
+	"github.com/volcano-sh/agentcube/pkg/redis"
 )
 
 // Server is the main structure for Router apiserver
@@ -16,8 +19,26 @@ type Server struct {
 	engine         *gin.Engine
 	httpServer     *http.Server
 	sessionManager SessionManager
-	redisManager   RedisManager
-	semaphore      chan struct{} // For limiting concurrent requests
+	redisClient    redis.Client
+	semaphore      chan struct{}   // For limiting concurrent requests
+	httpTransport  *http.Transport // Reusable HTTP transport for connection pooling
+}
+
+// makeRedisOptions creates redis options from environment variables
+func makeRedisOptions() (*redisv9.Options, error) {
+	redisAddr := os.Getenv("REDIS_ADDR")
+	if redisAddr == "" {
+		return nil, fmt.Errorf("missing env var REDIS_ADDR")
+	}
+	redisPassword := os.Getenv("REDIS_PASSWORD")
+	if redisPassword == "" {
+		return nil, fmt.Errorf("missing env var REDIS_PASSWORD")
+	}
+	redisOptions := &redisv9.Options{
+		Addr:     redisAddr,
+		Password: redisPassword,
+	}
+	return redisOptions, nil
 }
 
 // NewServer creates a new Router API server instance
@@ -43,11 +64,18 @@ func NewServer(config *Config) (*Server, error) {
 		config.SessionExpireDuration = 3600 // Default 1 hour
 	}
 
-	// Create session manager (using mock implementation)
-	sessionManager := NewMockSessionManager(config.SandboxEndpoints)
+	// Initialize Redis client
+	redisOptions, err := makeRedisOptions()
+	if err != nil {
+		return nil, fmt.Errorf("make redis options failed: %w", err)
+	}
+	redisClient := redis.NewClient(redisOptions)
 
-	// Create Redis manager (using mock implementation)
-	redisManager := NewMockRedisManager(config.EnableRedis)
+	// Create session manager with redis client
+	sessionManager, err := NewSessionManager(redisClient)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create session manager: %w", err)
+	}
 
 	// Set Gin mode based on environment
 	if config.Debug {
@@ -56,11 +84,21 @@ func NewServer(config *Config) (*Server, error) {
 		gin.SetMode(gin.ReleaseMode)
 	}
 
+	// Create a reusable HTTP transport for connection pooling
+	httpTransport := &http.Transport{
+		MaxIdleConns:        config.MaxIdleConns,
+		MaxIdleConnsPerHost: config.MaxConnsPerHost,
+		IdleConnTimeout:     90 * time.Second,
+		DisableCompression:  false,
+		ForceAttemptHTTP2:   true,
+	}
+
 	server := &Server{
 		config:         config,
 		sessionManager: sessionManager,
-		redisManager:   redisManager,
+		redisClient:    redisClient,
 		semaphore:      make(chan struct{}, config.MaxConcurrentRequests),
+		httpTransport:  httpTransport,
 	}
 
 	// Setup routes

pkg/router/apiserver_test.go

@@ -2,11 +2,22 @@ package router
 
 import (
 	"context"
+	"os"
 	"testing"
 	"time"
 )
 
 func TestNewServer(t *testing.T) {
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
+
 	tests := []struct {
 		name    string
 		config  *Config
@@ -76,9 +87,9 @@ func TestNewServer(t *testing.T) {
 					t.Error("Session manager was not created")
 				}
 
-				// Verify redis manager was created
-				if server.redisManager == nil {
-					t.Error("Redis manager was not created")
+				// Verify redis client was created
+				if server.redisClient == nil {
+					t.Error("Redis client was not created")
 				}
 
 				// Verify semaphore was created with correct capacity
@@ -112,6 +123,16 @@ func TestNewServer(t *testing.T) {
 }
 
 func TestServer_DefaultValues(t *testing.T) {
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
+
 	config := &Config{
 		Port: "8080",
 		// Leave other values as zero to test defaults
@@ -145,6 +166,16 @@ func TestServer_DefaultValues(t *testing.T) {
 }
 
 func TestServer_ConcurrencyLimitMiddleware(t *testing.T) {
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
+
 	config := &Config{
 		Port:                  "8080",
 		MaxConcurrentRequests: 2, // Small limit for testing
@@ -168,6 +199,16 @@ func TestServer_ConcurrencyLimitMiddleware(t *testing.T) {
 }
 
 func TestServer_SetupRoutes(t *testing.T) {
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
+
 	config := &Config{
 		Port: "8080",
 	}
@@ -187,6 +228,16 @@ func TestServer_SetupRoutes(t *testing.T) {
 }
 
 func TestServer_StartContext(t *testing.T) {
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
+
 	config := &Config{
 		Port: "0", // Use port 0 to let the OS assign a free port
 	}
@@ -225,6 +276,16 @@ func TestServer_StartContext(t *testing.T) {
 }
 
 func TestServer_TLSConfiguration(t *testing.T) {
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
+
 	tests := []struct {
 		name      string
 		config    *Config
@@ -294,50 +355,27 @@ func TestServer_TLSConfiguration(t *testing.T) {
 }
 
 func TestServer_RedisIntegration(t *testing.T) {
-	tests := []struct {
-		name        string
-		enableRedis bool
-	}{
-		{
-			name:        "Redis enabled",
-			enableRedis: true,
-		},
-		{
-			name:        "Redis disabled",
-			enableRedis: false,
-		},
-	}
-
-	for _, tt := range tests {
-		t.Run(tt.name, func(t *testing.T) {
-			config := &Config{
-				Port:        "8080",
-				EnableRedis: tt.enableRedis,
-			}
+	// Set required environment variables for tests
+	os.Setenv("REDIS_ADDR", "localhost:6379")
+	os.Setenv("REDIS_PASSWORD", "test-password")
+	os.Setenv("WORKLOAD_MGR_URL", "http://localhost:8080")
+	defer func() {
+		os.Unsetenv("REDIS_ADDR")
+		os.Unsetenv("REDIS_PASSWORD")
+		os.Unsetenv("WORKLOAD_MGR_URL")
+	}()
 
-			server, err := NewServer(config)
-			if err != nil {
-				t.Fatalf("Failed to create server: %v", err)
-			}
+	config := &Config{
+		Port: "8080",
+	}
 
-			// Verify Redis manager was created
-			if server.redisManager == nil {
-				t.Error("Redis manager was not created")
-			}
+	server, err := NewServer(config)
+	if err != nil {
+		t.Fatalf("Failed to create server: %v", err)
+	}
 
-			// Test Redis manager functionality based on enabled state
-			err = server.redisManager.UpdateSessionActivity("test-session")
-			if tt.enableRedis {
-				if err != nil {
-					t.Errorf("Expected no error when Redis is enabled, got: %v", err)
-				}
-			} else {
-				// When disabled, UpdateSessionActivity should not return an error
-				// (it silently skips)
-				if err != nil {
-					t.Errorf("Expected no error when Redis is disabled, got: %v", err)
-				}
-			}
-		})
+	// Verify Redis client was created
+	if server.redisClient == nil {
+		t.Error("Redis client was not created")
 	}
 }