Genarrative/apps/ai-game-creator-shell/src-tauri/src/resource_preview_scheduler.rs

use std::collections::{HashMap, HashSet, VecDeque};
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::{Arc, Mutex};

use tokio::sync::{watch, Semaphore};

pub(crate) const PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR: &str =
    "project-resource-preview-scope-cancelled";

const PROJECT_RESOURCE_PREVIEW_GLOBAL_CONCURRENCY: usize = 3;
const PROJECT_RESOURCE_PREVIEW_MAX_REGISTERED_REQUESTS: usize = 4_096;
const PROJECT_RESOURCE_PREVIEW_SEEN_REQUEST_LIMIT: usize = 8_192;
const PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT: usize = 1_024;
const PROJECT_RESOURCE_PREVIEW_OPAQUE_ID_MAX_BYTES: usize = 160;

#[derive(Debug)]
struct ProjectResourcePreviewScopeCancellationInner {
    cancelled: AtomicBool,
    signal: watch::Sender<bool>,
}

#[derive(Clone, Debug)]
pub(crate) struct ProjectResourcePreviewScopeCancellation {
    inner: Arc<ProjectResourcePreviewScopeCancellationInner>,
}

impl ProjectResourcePreviewScopeCancellation {
    fn new() -> Self {
        let (signal, _receiver) = watch::channel(false);
        Self {
            inner: Arc::new(ProjectResourcePreviewScopeCancellationInner {
                cancelled: AtomicBool::new(false),
                signal,
            }),
        }
    }

    pub(crate) fn uncancelled() -> Self {
        Self::new()
    }

    pub(crate) fn is_cancelled(&self) -> bool {
        self.inner.cancelled.load(Ordering::Acquire)
    }

    pub(crate) fn check(&self) -> Result<(), String> {
        if self.is_cancelled() {
            Err(PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR.to_string())
        } else {
            Ok(())
        }
    }

    pub(crate) fn cancel(&self) {
        if !self.inner.cancelled.swap(true, Ordering::AcqRel) {
            self.inner.signal.send_replace(true);
        }
    }

    async fn cancelled(&self) {
        let mut receiver = self.inner.signal.subscribe();
        loop {
            if self.is_cancelled() || *receiver.borrow() {
                return;
            }
            if receiver.changed().await.is_err() {
                return;
            }
        }
    }
}

#[derive(Debug)]
struct ProjectResourcePreviewScopeEntry {
    cancellation: ProjectResourcePreviewScopeCancellation,
    request_count: usize,
}

#[derive(Debug, Default)]
struct ProjectResourcePreviewRegistry {
    scopes: HashMap<String, ProjectResourcePreviewScopeEntry>,
    active_requests: HashMap<String, String>,
    seen_request_ids: HashSet<String>,
    seen_request_order: VecDeque<String>,
    cancelled_scope_ids: HashSet<String>,
    cancelled_scope_order: VecDeque<String>,
}

#[derive(Debug)]
struct ProjectResourcePreviewReadManagerInner {
    permits: Arc<Semaphore>,
    registry: Mutex<ProjectResourcePreviewRegistry>,
}

#[derive(Clone, Debug)]
pub(crate) struct ProjectResourcePreviewReadManager {
    inner: Arc<ProjectResourcePreviewReadManagerInner>,
}

impl Default for ProjectResourcePreviewReadManager {
    fn default() -> Self {
        Self::new(PROJECT_RESOURCE_PREVIEW_GLOBAL_CONCURRENCY)
    }
}

impl ProjectResourcePreviewReadManager {
    fn new(concurrency: usize) -> Self {
        Self {
            inner: Arc::new(ProjectResourcePreviewReadManagerInner {
                permits: Arc::new(Semaphore::new(concurrency)),
                registry: Mutex::new(ProjectResourcePreviewRegistry::default()),
            }),
        }
    }

    pub(crate) async fn run<T, F>(
        &self,
        scope_id: &str,
        request_id: &str,
        operation: F,
    ) -> Result<T, String>
    where
        T: Send + 'static,
        F: FnOnce(&ProjectResourcePreviewScopeCancellation) -> Result<T, String> + Send + 'static,
    {
        let registration = self.register(scope_id, request_id)?;
        let cancellation = registration.cancellation.clone();
        let permits = Arc::clone(&self.inner.permits);
        let permit = tokio::select! {
            biased;
            _ = cancellation.cancelled() => {
                return Err(PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR.to_string());
            }
            permit = permits.acquire_owned() => {
                permit.map_err(|_| "资源预览读取管理器已关闭".to_string())?
            }
        };
        cancellation.check()?;
        let blocking_cancellation = cancellation.clone();
        let result = tokio::task::spawn_blocking(move || {
            let _permit_guard = permit;
            let _registration_guard = registration;
            blocking_cancellation.check()?;
            let result = operation(&blocking_cancellation)?;
            blocking_cancellation.check()?;
            Ok(result)
        })
        .await
        .map_err(|_| "资源预览读取任务异常结束".to_string())?;
        result
    }

    pub(crate) fn cancel_scope(&self, scope_id: &str) -> Result<(), String> {
        let scope_id = validate_opaque_id("scopeId", scope_id)?;
        let cancellation = {
            let mut registry = self
                .inner
                .registry
                .lock()
                .map_err(|_| "资源预览读取登记表不可用".to_string())?;
            remember_cancelled_scope(&mut registry, scope_id);
            registry
                .scopes
                .get(scope_id)
                .map(|entry| entry.cancellation.clone())
        };
        if let Some(cancellation) = cancellation {
            cancellation.cancel();
        }
        Ok(())
    }

    fn register(
        &self,
        scope_id: &str,
        request_id: &str,
    ) -> Result<ProjectResourcePreviewRequestRegistration, String> {
        let scope_id = validate_opaque_id("scopeId", scope_id)?;
        let request_id = validate_opaque_id("requestId", request_id)?;
        let mut registry = self
            .inner
            .registry
            .lock()
            .map_err(|_| "资源预览读取登记表不可用".to_string())?;
        if registry.cancelled_scope_ids.contains(scope_id) {
            return Err(PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR.to_string());
        }
        if registry.seen_request_ids.contains(request_id) {
            return Err("资源预览 requestId 已经使用".to_string());
        }
        if registry.active_requests.len() >= PROJECT_RESOURCE_PREVIEW_MAX_REGISTERED_REQUESTS {
            return Err("资源预览等待队列已满".to_string());
        }
        remember_request_id(&mut registry, request_id);
        let cancellation = {
            let scope = registry
                .scopes
                .entry(scope_id.to_string())
                .or_insert_with(|| ProjectResourcePreviewScopeEntry {
                    cancellation: ProjectResourcePreviewScopeCancellation::new(),
                    request_count: 0,
                });
            scope.request_count += 1;
            scope.cancellation.clone()
        };
        registry
            .active_requests
            .insert(request_id.to_string(), scope_id.to_string());
        Ok(ProjectResourcePreviewRequestRegistration {
            manager: self.clone(),
            request_id: request_id.to_string(),
            cancellation,
        })
    }

    fn finish_request(&self, request_id: &str) {
        let Ok(mut registry) = self.inner.registry.lock() else {
            return;
        };
        let Some(scope_id) = registry.active_requests.remove(request_id) else {
            return;
        };
        let remove_scope = if let Some(scope) = registry.scopes.get_mut(&scope_id) {
            scope.request_count = scope.request_count.saturating_sub(1);
            scope.request_count == 0
        } else {
            false
        };
        if remove_scope {
            registry.scopes.remove(&scope_id);
            trim_cancelled_scopes(&mut registry);
        }
    }

    #[cfg(test)]
    fn active_registry_counts(&self) -> (usize, usize) {
        let registry = self.inner.registry.lock().expect("preview registry");
        (registry.active_requests.len(), registry.scopes.len())
    }
}

#[derive(Debug)]
struct ProjectResourcePreviewRequestRegistration {
    manager: ProjectResourcePreviewReadManager,
    request_id: String,
    cancellation: ProjectResourcePreviewScopeCancellation,
}

impl Drop for ProjectResourcePreviewRequestRegistration {
    fn drop(&mut self) {
        self.manager.finish_request(&self.request_id);
    }
}

fn validate_opaque_id<'a>(label: &str, value: &'a str) -> Result<&'a str, String> {
    let value = value.trim();
    if value.is_empty()
        || value.len() > PROJECT_RESOURCE_PREVIEW_OPAQUE_ID_MAX_BYTES
        || !value
            .bytes()
            .all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'-' | b'_' | b'.' | b':'))
    {
        return Err(format!(
            "资源预览 {label} 必须是长度不超过 {PROJECT_RESOURCE_PREVIEW_OPAQUE_ID_MAX_BYTES} 的不透明标识"
        ));
    }
    Ok(value)
}

fn remember_request_id(registry: &mut ProjectResourcePreviewRegistry, request_id: &str) {
    registry.seen_request_ids.insert(request_id.to_string());
    registry
        .seen_request_order
        .push_back(request_id.to_string());
    let scan_limit = registry.seen_request_order.len();
    for _ in 0..scan_limit {
        if registry.seen_request_order.len() <= PROJECT_RESOURCE_PREVIEW_SEEN_REQUEST_LIMIT {
            break;
        }
        let Some(oldest) = registry.seen_request_order.pop_front() else {
            break;
        };
        if registry.active_requests.contains_key(&oldest) {
            registry.seen_request_order.push_back(oldest);
            continue;
        }
        registry.seen_request_ids.remove(&oldest);
    }
}

fn remember_cancelled_scope(registry: &mut ProjectResourcePreviewRegistry, scope_id: &str) {
    if registry.cancelled_scope_ids.insert(scope_id.to_string()) {
        registry
            .cancelled_scope_order
            .push_back(scope_id.to_string());
    }
    trim_cancelled_scopes(registry);
}

fn trim_cancelled_scopes(registry: &mut ProjectResourcePreviewRegistry) {
    let scan_limit = registry.cancelled_scope_order.len();
    for _ in 0..scan_limit {
        if registry.cancelled_scope_order.len() <= PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT {
            break;
        }
        let Some(oldest) = registry.cancelled_scope_order.pop_front() else {
            break;
        };
        if registry.scopes.contains_key(&oldest) {
            registry.cancelled_scope_order.push_back(oldest);
            continue;
        }
        registry.cancelled_scope_ids.remove(&oldest);
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::atomic::{AtomicUsize, Ordering};
    use std::time::Duration;

    #[tokio::test(flavor = "multi_thread", worker_threads = 4)]
    async fn global_permits_bound_reads_and_scope_cancel_releases_new_work() {
        let manager = ProjectResourcePreviewReadManager::new(3);
        let active = Arc::new(AtomicUsize::new(0));
        let peak = Arc::new(AtomicUsize::new(0));
        let entered = Arc::new(AtomicUsize::new(0));
        let mut old_tasks = Vec::new();
        for index in 0..3 {
            let manager = manager.clone();
            let active = Arc::clone(&active);
            let peak = Arc::clone(&peak);
            let entered = Arc::clone(&entered);
            old_tasks.push(tokio::spawn(async move {
                manager
                    .run("scope-a", &format!("request-a-{index}"), move |cancel| {
                        let current = active.fetch_add(1, Ordering::SeqCst) + 1;
                        peak.fetch_max(current, Ordering::SeqCst);
                        entered.fetch_add(1, Ordering::SeqCst);
                        while cancel.check().is_ok() {
                            std::thread::yield_now();
                        }
                        active.fetch_sub(1, Ordering::SeqCst);
                        cancel.check()?;
                        Ok(())
                    })
                    .await
            }));
        }
        tokio::time::timeout(Duration::from_secs(2), async {
            while entered.load(Ordering::SeqCst) != 3 {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("old scope occupies all permits");

        let new_started = Arc::new(AtomicBool::new(false));
        let manager_for_new = manager.clone();
        let new_started_for_task = Arc::clone(&new_started);
        let new_task = tokio::spawn(async move {
            manager_for_new
                .run("scope-b", "request-b-0", move |_| {
                    new_started_for_task.store(true, Ordering::SeqCst);
                    Ok(())
                })
                .await
        });
        tokio::task::yield_now().await;
        assert!(!new_started.load(Ordering::SeqCst));

        manager.cancel_scope("scope-a").expect("cancel old scope");
        assert_eq!(new_task.await.expect("new task join"), Ok(()));
        for task in old_tasks {
            assert_eq!(
                task.await.expect("old task join"),
                Err(PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR.to_string())
            );
        }
        assert_eq!(peak.load(Ordering::SeqCst), 3);
        assert_eq!(active.load(Ordering::SeqCst), 0);
        assert_eq!(manager.active_registry_counts(), (0, 0));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn queued_cancel_does_not_enter_operation_and_unknown_cancel_closes_late_request() {
        let manager = ProjectResourcePreviewReadManager::new(1);
        let blocker_entered = Arc::new(AtomicBool::new(false));
        let release_blocker = Arc::new(AtomicBool::new(false));
        let manager_for_blocker = manager.clone();
        let blocker_entered_for_task = Arc::clone(&blocker_entered);
        let release_blocker_for_task = Arc::clone(&release_blocker);
        let blocker = tokio::spawn(async move {
            manager_for_blocker
                .run("scope-blocker", "request-blocker", move |_| {
                    blocker_entered_for_task.store(true, Ordering::SeqCst);
                    while !release_blocker_for_task.load(Ordering::SeqCst) {
                        std::thread::yield_now();
                    }
                    Ok(())
                })
                .await
        });
        tokio::time::timeout(Duration::from_secs(2), async {
            while !blocker_entered.load(Ordering::SeqCst) {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("blocker starts");

        let queued_entered = Arc::new(AtomicBool::new(false));
        let manager_for_queued = manager.clone();
        let queued_entered_for_task = Arc::clone(&queued_entered);
        let queued = tokio::spawn(async move {
            manager_for_queued
                .run("scope-queued", "request-queued", move |_| {
                    queued_entered_for_task.store(true, Ordering::SeqCst);
                    Ok(())
                })
                .await
        });
        tokio::task::yield_now().await;
        manager
            .cancel_scope("scope-queued")
            .expect("cancel queued scope");
        assert_eq!(
            queued.await.expect("queued task join"),
            Err(PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR.to_string())
        );
        assert!(!queued_entered.load(Ordering::SeqCst));

        manager
            .cancel_scope("scope-before-register")
            .expect("cancel unknown scope");
        let late_result = manager
            .run("scope-before-register", "request-late", |_| Ok(()))
            .await;
        assert_eq!(
            late_result,
            Err(PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR.to_string())
        );

        release_blocker.store(true, Ordering::SeqCst);
        assert_eq!(blocker.await.expect("blocker join"), Ok(()));
        assert_eq!(manager.active_registry_counts(), (0, 0));
    }

    #[tokio::test]
    async fn duplicate_request_id_is_rejected_after_completion() {
        let manager = ProjectResourcePreviewReadManager::new(1);
        assert_eq!(
            manager.run("scope-a", "same-request", |_| Ok(())).await,
            Ok(())
        );
        let duplicate = manager
            .run("scope-b", "same-request", |_| Ok(()))
            .await
            .expect_err("duplicate request rejected");
        assert!(duplicate.contains("requestId 已经使用"));
    }

    #[tokio::test]
    async fn operation_error_and_panic_release_registration_and_permit() {
        let manager = ProjectResourcePreviewReadManager::new(1);
        assert_eq!(
            manager
                .run("scope-error", "request-error", |_| {
                    Err::<(), _>("预览读取失败".to_string())
                })
                .await,
            Err("预览读取失败".to_string())
        );
        assert_eq!(manager.active_registry_counts(), (0, 0));

        let panic_error = manager
            .run("scope-panic", "request-panic", |_| -> Result<(), String> {
                panic!("preview worker panic")
            })
            .await
            .expect_err("panic is converted into a stable command error");
        assert!(panic_error.contains("资源预览读取任务异常结束"));
        assert_eq!(manager.active_registry_counts(), (0, 0));

        assert_eq!(
            manager
                .run("scope-after-panic", "request-after-panic", |_| Ok(()))
                .await,
            Ok(())
        );
        assert_eq!(manager.active_registry_counts(), (0, 0));
    }

    #[tokio::test(flavor = "multi_thread", worker_threads = 2)]
    async fn aborting_the_outer_future_keeps_the_blocking_read_registered_and_permitted() {
        let manager = ProjectResourcePreviewReadManager::new(1);
        let entered = Arc::new(AtomicBool::new(false));
        let cancellation_observed = Arc::new(AtomicBool::new(false));
        let (release_sender, release_receiver) = std::sync::mpsc::sync_channel::<()>(0);
        let manager_for_old = manager.clone();
        let entered_for_old = Arc::clone(&entered);
        let cancellation_observed_for_old = Arc::clone(&cancellation_observed);
        let old_task = tokio::spawn(async move {
            manager_for_old
                .run("scope-aborted", "request-aborted", move |cancel| {
                    entered_for_old.store(true, Ordering::SeqCst);
                    loop {
                        if let Err(error) = cancel.check() {
                            cancellation_observed_for_old.store(true, Ordering::SeqCst);
                            return Err(error);
                        }
                        match release_receiver.recv_timeout(Duration::from_millis(5)) {
                            Ok(()) | Err(std::sync::mpsc::RecvTimeoutError::Disconnected) => {
                                return Ok(())
                            }
                            Err(std::sync::mpsc::RecvTimeoutError::Timeout) => {}
                        }
                    }
                })
                .await
        });
        tokio::time::timeout(Duration::from_secs(2), async {
            while !entered.load(Ordering::SeqCst) {
                tokio::task::yield_now().await;
            }
        })
        .await
        .expect("blocking read starts");

        old_task.abort();
        assert!(old_task
            .await
            .expect_err("outer future is aborted")
            .is_cancelled());
        assert_eq!(manager.active_registry_counts(), (1, 1));

        let new_started = Arc::new(AtomicBool::new(false));
        let manager_for_new = manager.clone();
        let new_started_for_task = Arc::clone(&new_started);
        let new_task = tokio::spawn(async move {
            manager_for_new
                .run("scope-new", "request-new", move |_| {
                    new_started_for_task.store(true, Ordering::SeqCst);
                    Ok(())
                })
                .await
        });
        tokio::task::yield_now().await;
        assert!(!new_started.load(Ordering::SeqCst));

        manager
            .cancel_scope("scope-aborted")
            .expect("cancel blocking read after outer abort");
        if tokio::time::timeout(Duration::from_secs(2), async {
            while !cancellation_observed.load(Ordering::SeqCst) {
                tokio::task::yield_now().await;
            }
        })
        .await
        .is_err()
        {
            drop(release_sender);
            panic!("blocking read did not observe scope cancellation after outer abort");
        }
        drop(release_sender);
        assert_eq!(
            tokio::time::timeout(Duration::from_secs(2), new_task)
                .await
                .expect("new task starts after cancelled blocking read exits")
                .expect("new task join"),
            Ok(())
        );
        assert!(new_started.load(Ordering::SeqCst));
        assert_eq!(manager.active_registry_counts(), (0, 0));
    }

    #[test]
    fn seen_request_tombstones_evict_the_oldest_entry_at_the_limit() {
        let mut registry = ProjectResourcePreviewRegistry::default();
        for index in 0..=PROJECT_RESOURCE_PREVIEW_SEEN_REQUEST_LIMIT {
            remember_request_id(&mut registry, &format!("request-{index}"));
        }

        assert_eq!(
            registry.seen_request_ids.len(),
            PROJECT_RESOURCE_PREVIEW_SEEN_REQUEST_LIMIT
        );
        assert_eq!(
            registry.seen_request_order.len(),
            PROJECT_RESOURCE_PREVIEW_SEEN_REQUEST_LIMIT
        );
        assert!(!registry.seen_request_ids.contains("request-0"));
        assert!(registry.seen_request_ids.contains("request-1"));
        assert!(registry.seen_request_ids.contains(&format!(
            "request-{}",
            PROJECT_RESOURCE_PREVIEW_SEEN_REQUEST_LIMIT
        )));
    }

    #[test]
    fn cancelled_scope_tombstones_evict_the_oldest_entry_at_the_limit() {
        let mut registry = ProjectResourcePreviewRegistry::default();
        for index in 0..=PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT {
            remember_cancelled_scope(&mut registry, &format!("scope-{index}"));
        }

        assert_eq!(
            registry.cancelled_scope_ids.len(),
            PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT
        );
        assert_eq!(
            registry.cancelled_scope_order.len(),
            PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT
        );
        assert!(!registry.cancelled_scope_ids.contains("scope-0"));
        assert!(registry.cancelled_scope_ids.contains("scope-1"));
        assert!(registry.cancelled_scope_ids.contains(&format!(
            "scope-{}",
            PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT
        )));
    }

    #[test]
    fn active_cancelled_scope_tombstones_converge_after_the_scope_finishes() {
        let manager = ProjectResourcePreviewReadManager::new(1);
        let mut registrations = Vec::new();
        for index in 0..=PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT {
            let scope_id = format!("active-scope-{index}");
            let request_id = format!("active-request-{index}");
            registrations.push(
                manager
                    .register(&scope_id, &request_id)
                    .expect("register active scope"),
            );
            manager
                .cancel_scope(&scope_id)
                .expect("cancel active scope");
        }

        {
            let registry = manager.inner.registry.lock().expect("preview registry");
            assert_eq!(
                registry.cancelled_scope_ids.len(),
                PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT + 1
            );
            assert_eq!(
                registry.cancelled_scope_order.len(),
                PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT + 1
            );
        }
        assert_eq!(
            manager
                .register("active-scope-0", "late-request")
                .expect_err("active cancelled scope stays closed"),
            PROJECT_RESOURCE_PREVIEW_CANCELLED_ERROR
        );

        let completed_scope = format!(
            "active-scope-{}",
            PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT
        );
        drop(registrations.pop());
        let registry = manager.inner.registry.lock().expect("preview registry");
        assert_eq!(
            registry.cancelled_scope_ids.len(),
            PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT
        );
        assert_eq!(
            registry.cancelled_scope_order.len(),
            PROJECT_RESOURCE_PREVIEW_CANCELLED_SCOPE_LIMIT
        );
        assert!(!registry.cancelled_scope_ids.contains(&completed_scope));
        assert!(registry.cancelled_scope_ids.contains("active-scope-0"));
    }
}