linkrs_server/session/

session_manager.rs

use dashmap::DashMap;
use linkrs_core::event_dispatch::{EventFilter, EventSubscriptions, SubscriptionId};
use linkrs_query::{SessionEvent, SessionEventCallback};
use log::{info, warn};
use parking_lot::RwLock;
use rand::RngCore;
use std::collections::HashMap;
use std::sync::{
    atomic::{AtomicBool, AtomicU64, Ordering},
    Arc,
};
use std::time::{Duration, Instant, SystemTime};
use tokio::time;

use super::{SessionError, SessionResult};
use crate::client::{ClientSession, Session};

pub const DEFAULT_MAX_ALLOWED_CONNECTIONS: usize = 100; // Default maximum number of connections (in a single-node scenario)
pub const DEFAULT_SESSION_IDLE_TIMEOUT: Duration = Duration::from_secs(600); // 10 minutes

/// Session information, used to display the list of sessions
#[derive(Debug, Clone)]
pub struct SessionInfo {
    pub session_id: i64,
    pub user_name: String,
    pub space_name: Option<String>,
    pub graph_addr: Option<String>,
    pub create_time: SystemTime,
    pub last_access_time: SystemTime,
    pub active_queries: usize,
    pub timezone: Option<i32>,
}

impl SessionInfo {
    pub fn from_client_session(session: &ClientSession, create_time: SystemTime) -> Self {
        Self {
            session_id: session.id(),
            user_name: session.user(),
            space_name: session.space_name(),
            graph_addr: session.graph_addr(),
            create_time,
            last_access_time: SystemTime::now() - Duration::from_millis(session.idle_seconds()),
            active_queries: session.active_queries_count(),
            timezone: session.timezone(),
        }
    }

    pub fn from_params(
        session_id_str: &str,
        user_name: &str,
        space_name: Option<String>,
        client_ip: &str,
        client_port: u16,
    ) -> Result<Self, String> {
        let session_id = session_id_str
            .parse::<i64>()
            .map_err(|_| format!("Invalid session ID: {}", session_id_str))?;

        let graph_addr = if client_ip.is_empty() {
            None
        } else {
            Some(format!("{}:{}", client_ip, client_port))
        };

        let now = SystemTime::now();
        Ok(Self {
            session_id,
            user_name: user_name.to_string(),
            space_name,
            graph_addr,
            create_time: now,
            last_access_time: now,
            active_queries: 0,
            timezone: None,
        })
    }

    pub fn touch(&mut self) {
        self.last_access_time = SystemTime::now();
    }
}

pub struct GraphSessionManager {
    // Use DashMap to achieve true concurrent access without the need for explicit locking.
    sessions: Arc<DashMap<i64, Arc<ClientSession>>>,
    active_sessions: Arc<DashMap<i64, Instant>>, // session_id -> last_activity_time
    // Read more, write less – and use RwLock.
    session_create_times: Arc<RwLock<HashMap<i64, SystemTime>>>, // session_id -> create_time
    host_addr: String,
    max_connections: usize,
    session_idle_timeout: Duration,
    /// Initial auto-commit mode for newly created sessions, from
    /// `[transaction].auto_commit`. Explicit BEGIN/COMMIT still overrides it.
    default_auto_commit: AtomicBool,
    /// Is the background cleanup task currently running?
    cleanup_task_running: Arc<AtomicBool>,
    /// Cumulative number of sessions created (monotonic, never decremented).
    total_sessions_created: Arc<AtomicU64>,
    // Shared with `QueryManager` when co-located: single enum dual emission
    // source (`SessionCreated/Destroyed` here, `Query*` there).
    session_callbacks: Arc<EventSubscriptions<SessionEvent>>,
}

impl std::fmt::Debug for GraphSessionManager {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.debug_struct("GraphSessionManager")
            .field("sessions_count", &self.sessions.len())
            .field("host_addr", &self.host_addr)
            .field("max_connections", &self.max_connections)
            .field("session_callbacks", &self.session_callbacks.len())
            .finish()
    }
}

impl GraphSessionManager {
    /// Create a new session manager.
    ///
    /// This constructor does not automatically initiate background cleanup tasks.
    /// Need to explicitly call `start_cleanup_task()` to start
    pub fn new(
        host_addr: String,
        max_connections: usize,
        session_idle_timeout: Duration,
    ) -> Arc<Self> {
        Self::new_with_shared(
            host_addr,
            max_connections,
            session_idle_timeout,
            Arc::new(EventSubscriptions::new()),
        )
    }

    /// Build a manager sharing one session-event registry with `QueryManager`.
    pub fn new_with_shared(
        host_addr: String,
        max_connections: usize,
        session_idle_timeout: Duration,
        shared: Arc<EventSubscriptions<SessionEvent>>,
    ) -> Arc<Self> {
        Arc::new(Self {
            sessions: Arc::new(DashMap::new()),
            active_sessions: Arc::new(DashMap::new()),
            session_create_times: Arc::new(RwLock::new(HashMap::new())),
            host_addr,
            max_connections,
            session_idle_timeout,
            default_auto_commit: AtomicBool::new(true),
            cleanup_task_running: Arc::new(AtomicBool::new(false)),
            total_sessions_created: Arc::new(AtomicU64::new(0)),
            session_callbacks: shared,
        })
    }

    /// Shared session-event registry behind this manager.
    pub fn shared_session_callbacks(&self) -> Arc<EventSubscriptions<SessionEvent>> {
        Arc::clone(&self.session_callbacks)
    }

    /// Mask a session id for logs: prefix plus a short hash only.
    pub fn mask_session_id(session_id: i64) -> String {
        use std::collections::hash_map::DefaultHasher;
        use std::hash::{Hash, Hasher};
        let mut hasher = DefaultHasher::new();
        session_id.hash(&mut hasher);
        format!("{:x}..{:04x}", session_id >> 48, hasher.finish() & 0xffff)
    }

    /// Session idle timeout driving reclamation and expiry responses.
    pub fn idle_timeout(&self) -> Duration {
        self.session_idle_timeout
    }

    /// Set the auto-commit mode applied to sessions created from now on.
    pub fn set_default_auto_commit(&self, auto_commit: bool) {
        self.default_auto_commit
            .store(auto_commit, Ordering::SeqCst);
    }

    /// Point this manager at a shared session-event registry.
    ///
    /// The field is an `Arc`, so an already-published instance needs `&mut`
    /// access to switch sources; managers behind `Arc` should prefer
    /// `new_with_shared` at construction time instead.
    pub fn set_shared_session_callbacks(&mut self, shared: Arc<EventSubscriptions<SessionEvent>>) {
        self.session_callbacks = shared;
    }

    /// Register a runtime observer for session lifecycle events.
    pub fn register_session_callback(&self, callback: SessionEventCallback) -> SubscriptionId {
        self.session_callbacks.add(callback)
    }

    /// Register a filtered observer invoked only when `filter` returns true.
    pub fn register_session_callback_filtered(
        &self,
        callback: SessionEventCallback,
        filter: EventFilter<SessionEvent>,
    ) -> SubscriptionId {
        self.session_callbacks.add_filtered(callback, Some(filter))
    }

    /// Remove a previously registered observer. Returns true if present.
    pub fn unregister_session_callback(&self, id: SubscriptionId) -> bool {
        self.session_callbacks.remove(id)
    }

    /// Number of registered session observers.
    pub fn session_callback_count(&self) -> usize {
        self.session_callbacks.len()
    }

    fn emit_session_event(&self, event: SessionEvent) {
        if self.session_callbacks.is_empty() {
            return;
        }
        self.session_callbacks.dispatch("session", &event);
    }

    /// Start the background session cleanup task.
    ///
    /// If the task is already running, this method will not start it again.
    pub async fn start_cleanup_task(self: &Arc<Self>) {
        if self.cleanup_task_running.swap(true, Ordering::SeqCst) {
            info!("Session cleanup task is already running");
            return;
        }

        info!("Starting session cleanup task");
        let manager_clone = Arc::clone(self);
        tokio::spawn(async move {
            manager_clone.background_reclamation_task().await;
        });
    }

    /// Stop the background session cleaning task.
    ///
    /// Set a stop flag; the background task will exit during the next iteration.
    pub fn stop_cleanup_task(&self) {
        info!("Stopping session cleanup task");
        self.cleanup_task_running.store(false, Ordering::SeqCst);
    }

    /// Check whether the background cleanup task is currently running.
    pub fn is_cleanup_task_running(&self) -> bool {
        self.cleanup_task_running.load(Ordering::SeqCst)
    }

    /// Creates a new session
    pub async fn create_session(
        &self,
        user_name: String,
        _client_ip: String,
    ) -> Result<Arc<ClientSession>, String> {
        info!("Creating new session for user: {}", user_name);

        // Check if we're out of connections
        if self.is_out_of_connections().await {
            warn!(
                "Failed to create session for user {}: maximum connections exceeded",
                user_name
            );
            return Err("Exceeded maximum allowed connections".to_string());
        }

        // Generate a new session ID, retrying on the unlikely event of a
        // collision with a live session.
        let session_id = loop {
            let candidate = self.generate_session_id();
            if !self.sessions.contains_key(&candidate) {
                break candidate;
            }
        };
        info!(
            "Generated session {} for user: {}",
            Self::mask_session_id(session_id),
            user_name
        );

        let session = Session {
            session_id,
            user_name: user_name.clone(),
            space_name: None,
            graph_addr: Some(self.host_addr.clone()),
            timezone: None,
        };

        let client_session = ClientSession::new(session);
        client_session.set_auto_commit(self.default_auto_commit.load(Ordering::SeqCst));

        // Add to sessions and active sessions
        let create_time = SystemTime::now();

        // DashMap does not require explicit locking; concurrent insertions are possible without any issues.
        self.sessions
            .insert(session_id, Arc::clone(&client_session));
        self.active_sessions.insert(session_id, Instant::now());
        self.total_sessions_created.fetch_add(1, Ordering::Relaxed);

        // Write lock protection creation time
        {
            let mut create_times = self.session_create_times.write();
            create_times.insert(session_id, create_time);
        }

        info!(
            "Successfully created session {} for user: {}",
            Self::mask_session_id(session_id),
            user_name
        );
        self.emit_session_event(SessionEvent::SessionCreated {
            session_id,
            user_name,
        });
        Ok(client_session)
    }

    /// Finds an existing session
    pub fn find_session(&self, session_id: i64) -> Option<Arc<ClientSession>> {
        // DashMap supports true concurrent reading without the need for locking.
        self.sessions.get(&session_id).map(|entry| entry.clone())
    }

    /// Finds an existing session only from local cache
    pub fn find_session_from_cache(&self, session_id: i64) -> Option<Arc<ClientSession>> {
        self.find_session(session_id)
    }

    /// Removes a session from local cache
    pub async fn remove_session(&self, session_id: i64) {
        info!("Removing session {}", Self::mask_session_id(session_id));

        if let Some(entry) = self.sessions.get(&session_id) {
            entry.value().mark_all_queries_killed();
            entry.value().close_all_cursors();
        }
        // DashMap does not require explicit locking.
        let existed = self.sessions.remove(&session_id).is_some();
        self.active_sessions.remove(&session_id);

        // Write lock protection creation time
        {
            let mut create_times = self.session_create_times.write();
            create_times.remove(&session_id);
        }

        info!(
            "Successfully removed session {}",
            Self::mask_session_id(session_id)
        );
        if existed {
            self.emit_session_event(SessionEvent::SessionDestroyed { session_id });
        }
    }

    /// Invalidate all sessions of a user, optionally keeping one session.
    ///
    /// Used when a user is disabled or dropped so existing sessions stop
    /// working immediately. Returns the removed session ids.
    pub async fn remove_sessions_by_username(
        &self,
        username: &str,
        except_session_id: Option<i64>,
    ) -> Vec<i64> {
        let targets: Vec<i64> = self
            .sessions
            .iter()
            .filter(|entry| {
                entry.value().user() == username && Some(*entry.key()) != except_session_id
            })
            .map(|entry| *entry.key())
            .collect();
        for session_id in &targets {
            self.remove_session(*session_id).await;
        }
        targets
    }

    /// Gets all sessions from the local cache
    pub fn get_sessions_from_local_cache(&self) -> Vec<Session> {
        // DashMap supports iterators, and no locking is required.
        self.sessions
            .iter()
            .map(|entry| entry.value().get_session())
            .collect()
    }

    /// Obtain information about the session list, which is used for the SHOW SESSIONS command.
    pub async fn list_sessions(&self) -> Vec<SessionInfo> {
        // Read lock acquisition time
        let create_times = self.session_create_times.read();

        // DashMap’s iteration process does not require locking.
        self.sessions
            .iter()
            .filter_map(|entry| {
                let session_id = entry.key();
                let client_session = entry.value();
                create_times.get(session_id).map(|&create_time| {
                    SessionInfo::from_client_session(client_session, create_time)
                })
            })
            .collect()
    }

    /// Obtain detailed information about the specified session.
    pub async fn get_session_info(&self, session_id: i64) -> Option<SessionInfo> {
        // Reading DashMap does not require any locking mechanisms.
        let client_session = self.sessions.get(&session_id)?;

        // Read lock acquisition time
        let create_times = self.session_create_times.read();
        create_times
            .get(&session_id)
            .map(|&create_time| SessionInfo::from_client_session(&client_session, create_time))
    }

    /// Terminate the specified session (KILL SESSION)
    ///
    /// # Parameters
    /// `session_id` – The ID of the session that needs to be terminated.
    /// `current_user` – The username of the user who is performing the termination operation.
    /// `is_admin` – Determines whether the current user has the Admin role.
    ///
    /// # Return
    /// * `Ok(())` - Successfully terminate the session
    /// * `Err(SessionError)` - Specific reasons for termination failure
    pub async fn kill_session(
        &self,
        session_id: i64,
        current_user: &str,
        is_admin: bool,
    ) -> SessionResult<()> {
        info!(
            "Attempting to kill session {} by user: {} (is_admin: {})",
            Self::mask_session_id(session_id),
            current_user,
            is_admin
        );

        // Find the target conversation.
        let target_session = self
            .find_session(session_id)
            .ok_or(SessionError::session_not_found(session_id))?;

        let target_user = target_session.user();

        // Permission check: You can only terminate your own session, or you must have Admin privileges.
        if !is_admin && target_user != current_user {
            warn!(
                "User {} attempted to kill session {} without permission (target user: {})",
                current_user, session_id, target_user
            );
            return Err(SessionError::insufficient_permission());
        }

        info!(
            "Killing session {} (user: {}, active queries: {})",
            Self::mask_session_id(session_id),
            target_user,
            target_session.active_queries_count()
        );

        // Remove the session from the manager, which terminates all
        // queries and closes all cursors before dropping the entry.
        self.remove_session(session_id).await;

        info!(
            "Successfully killed session {} by user: {}",
            Self::mask_session_id(session_id),
            current_user
        );
        Ok(())
    }

    /// Batch termination of multiple sessions
    pub async fn kill_multiple_sessions(
        &self,
        session_ids: &[i64],
        current_user: &str,
        is_admin: bool,
    ) -> Vec<SessionResult<()>> {
        let mut results = Vec::with_capacity(session_ids.len());
        for &session_id in session_ids {
            results.push(self.kill_session(session_id, current_user, is_admin).await);
        }
        results
    }

    /// Whether exceeds the max allowed connections
    pub async fn is_out_of_connections(&self) -> bool {
        // DashMap's len() is an O(1) operation, no lock required
        self.active_sessions.len() >= self.max_connections
    }

    /// Obtain the number of active sessions
    pub async fn active_session_count(&self) -> usize {
        self.active_sessions.len()
    }

    /// Cumulative number of sessions created since process start.
    pub fn total_sessions_created(&self) -> u64 {
        self.total_sessions_created.load(Ordering::Relaxed)
    }

    /// Obtaining the maximum limit for the number of connections
    pub fn max_connections(&self) -> usize {
        self.max_connections
    }

    /// Generate a cryptographically random positive session ID.
    ///
    /// High-entropy 63-bit random values make session IDs unguessable even on
    /// public deployments, eliminating session-hijacking via brute force.
    fn generate_session_id(&self) -> i64 {
        let mut rng = rand::thread_rng();
        loop {
            let raw = rng.next_u64();
            // Keep the sign bit clear so we always emit positive i64 values.
            let candidate = (raw & 0x7FFF_FFFF_FFFF_FFFF) as i64;
            if candidate > 0 {
                return candidate;
            }
        }
    }

    /// Background task: Regularly clean up expired sessions.
    ///
    /// Check every 30 seconds and remove sessions that have exceeded the idle timeout period.
    /// Can be stopped via the `stop_cleanup_task()` method
    async fn background_reclamation_task(self: Arc<Self>) {
        let mut interval = time::interval(Duration::from_secs(30));

        loop {
            interval.tick().await;

            // Check whether it should be stopped.
            if !self.cleanup_task_running.load(Ordering::SeqCst) {
                info!("Session cleanup task is stopping");
                break;
            }

            self.reclaim_expired_sessions().await;
        }

        info!("Session cleanup task has stopped");
    }

    /// Reclaims expired sessions
    async fn reclaim_expired_sessions(&self) {
        // DashMap supports iteration without the need for locking.
        let expired_sessions: Vec<i64> = self
            .active_sessions
            .iter()
            .filter(|entry| entry.value().elapsed() > self.session_idle_timeout)
            .map(|entry| *entry.key())
            .collect();

        if !expired_sessions.is_empty() {
            info!(
                "Found {} expired sessions to reclaim",
                expired_sessions.len()
            );
        }

        // Remove expired sessions
        for session_id in expired_sessions {
            info!(
                "Reclaiming expired session {}",
                Self::mask_session_id(session_id)
            );
            self.remove_session(session_id).await;
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    fn create_test_session_manager() -> Arc<GraphSessionManager> {
        GraphSessionManager::new(
            "127.0.0.1:9669".to_string(),
            DEFAULT_MAX_ALLOWED_CONNECTIONS,
            DEFAULT_SESSION_IDLE_TIMEOUT,
        )
    }

    #[tokio::test]
    async fn test_session_manager_creation() {
        let session_manager = create_test_session_manager();

        assert_eq!(session_manager.host_addr, "127.0.0.1:9669");
        assert_eq!(session_manager.get_sessions_from_local_cache().len(), 0);
        assert!(!session_manager.is_cleanup_task_running());
    }

    #[tokio::test]
    async fn test_default_auto_commit_applies_to_new_sessions() {
        let session_manager = create_test_session_manager();

        let session = session_manager
            .create_session("autocommit".into(), "127.0.0.1".to_string())
            .await
            .expect("Failed to create session");
        assert!(session.is_auto_commit());

        session_manager.set_default_auto_commit(false);
        let session = session_manager
            .create_session("explicit".into(), "127.0.0.1".to_string())
            .await
            .expect("Failed to create session");
        assert!(!session.is_auto_commit());
    }

    #[tokio::test]
    async fn test_create_and_find_session() {
        let session_manager = create_test_session_manager();

        let session = session_manager
            .create_session("testuser".into(), "127.0.0.1".to_string())
            .await
            .expect("Failed to create session");

        assert_eq!(session.user(), "testuser");
        assert!(!session_manager.is_out_of_connections().await);

        let found_session = session_manager
            .find_session(session.id())
            .expect("Failed to find session");
        assert_eq!(found_session.user(), "testuser");

        // Test find non-existent session
        assert!(session_manager.find_session(999999).is_none());
    }

    #[tokio::test]
    async fn test_remove_session() {
        let session_manager = create_test_session_manager();

        let session = session_manager
            .create_session("testuser".into(), "127.0.0.1".to_string())
            .await
            .expect("Failed to create session");

        assert!(session_manager.find_session(session.id()).is_some());

        session_manager.remove_session(session.id()).await;
        assert!(session_manager.find_session(session.id()).is_none());
    }

    #[tokio::test]
    async fn test_max_connections() {
        let session_manager = GraphSessionManager::new(
            "127.0.0.1:9669".to_string(),
            5,
            DEFAULT_SESSION_IDLE_TIMEOUT,
        );

        assert!(!session_manager.is_out_of_connections().await);

        for i in 0..5 {
            let _ = session_manager
                .create_session(format!("user{}", i), "127.0.0.1".to_string())
                .await;
        }

        assert!(session_manager.is_out_of_connections().await);

        // Attempts to create a 6th session should fail
        let result = session_manager
            .create_session("user6".into(), "127.0.0.1".to_string())
            .await;
        assert!(result.is_err());
    }

    #[tokio::test]
    async fn test_kill_session() {
        let session_manager = create_test_session_manager();

        let session = session_manager
            .create_session("testuser".into(), "127.0.0.1".to_string())
            .await
            .expect("Failed to create session");

        let session_id = session.id();

        // Normal users try to terminate their own session - should succeed
        let result = session_manager
            .kill_session(session_id, "testuser", false)
            .await;
        assert!(result.is_ok());
        assert!(session_manager.find_session(session_id).is_none());

        // Create a new session to test permission checking
        let session2 = session_manager
            .create_session("user2".into(), "127.0.0.1".to_string())
            .await
            .expect("Failed to create session");

        // Normal user attempts to terminate another user's session - should fail
        let result = session_manager
            .kill_session(session2.id(), "otheruser", false)
            .await;
        assert!(result.is_err());

        // Admin Terminate another user's session - should succeed!
        let result = session_manager
            .kill_session(session2.id(), "admin", true)
            .await;
        assert!(result.is_ok());
    }

    #[tokio::test]
    async fn test_list_sessions() {
        let session_manager = create_test_session_manager();

        // Creating Multiple Sessions
        for i in 0..3 {
            let _ = session_manager
                .create_session(format!("user{}", i), "127.0.0.1".to_string())
                .await;
        }

        let sessions = session_manager.list_sessions().await;
        assert_eq!(sessions.len(), 3);
    }

    #[tokio::test]
    async fn test_concurrent_access() {
        use tokio::task;

        let session_manager = create_test_session_manager();
        let mut handles = vec![];

        // Concurrent Session Creation
        for i in 0..10 {
            let manager = Arc::clone(&session_manager);
            let handle = task::spawn(async move {
                manager
                    .create_session(format!("user{}", i), "127.0.0.1".to_string())
                    .await
            });
            handles.push(handle);
        }

        // Waiting for all tasks to be completed
        for handle in handles {
            let _ = handle.await.expect("Task should complete");
        }

        // Verify that all sessions were created successfully
        assert_eq!(session_manager.get_sessions_from_local_cache().len(), 10);
    }
}