core/test/rust/runtime/database/stream.rs

use super::*;
use crate::{
    adapter::Document,
    database::{Binding, Settings, Storage},
    task::{Limits, Runtime},
};
use std::{
    path::PathBuf,
    sync::atomic::{AtomicU64, AtomicUsize, Ordering},
};

static NEXT: AtomicU64 = AtomicU64::new(0);
struct Directory(PathBuf);
impl Directory {
    fn new() -> Self {
        let path = std::env::temp_dir().join(format!(
            "dever-row-stream-{}-{}",
            std::process::id(),
            NEXT.fetch_add(1, Ordering::Relaxed)
        ));
        std::fs::create_dir(&path).unwrap();
        Self(path)
    }
}
impl Drop for Directory {
    fn drop(&mut self) {
        std::fs::remove_dir_all(&self.0).unwrap();
    }
}

async fn registry(directory: &Directory, scope: &Scope<'_>) -> Arc<Registry> {
    let settings = Settings::load(&Document::parse(r#"{"database":{"default":{"type":"sqlite","path":"case.db","max_connections":1,"max_page_size":2,"wait_timeout_ms":1000}}}"#.into()).unwrap()).unwrap();
    Registry::prepare(
        &directory.0,
        settings,
        &[Binding {
            component: "demo",
            database: None,
            table: "rows",
            schema: Sql::same("CREATE TABLE rows (id INTEGER PRIMARY KEY)"),
            plan: None,
            references: &[],
            storage: Storage::Global,
        }],
        scope,
    )
    .await
    .unwrap()
}

const ROWS: Sql = Sql::same(
    "WITH RECURSIVE items(value) AS (SELECT 1 UNION ALL SELECT value+1 FROM items WHERE value < 10000) SELECT value FROM items",
);

#[test]
fn bounded_stream_close_drop_and_decode_error_release_single_connection() {
    let directory = Directory::new();
    let runtime = Runtime::new(Limits {
        threads: 2,
        workers: 1,
        tasks: 16,
    })
    .unwrap();
    runtime
        .run(async |scope| {
            let registry = registry(&directory, scope).await;
            let mut context = Context::new(registry.clone());
            for mode in 0..3 {
                let produced = Arc::new(AtomicUsize::new(0));
                let counted = produced.clone();
                let rows = context
                    .stream(0, scope, ROWS, vec![], move |row| {
                        counted.fetch_add(1, Ordering::Relaxed);
                        row.get::<_, i64>(0)
                    })
                    .await
                    .unwrap();
                assert_eq!(rows.next().await.unwrap(), Some(1));
                // 消费者暂停时最多只多解码通道一行与待发送一行。
                tokio::time::sleep(std::time::Duration::from_millis(10)).await;
                assert!(produced.load(Ordering::Relaxed) <= 3);
                // SQLite 游标在阻塞线程中借用 statement；背压期间仍持有 worker。
                assert!(
                    tokio::time::timeout(
                        std::time::Duration::from_millis(10),
                        scope.blocking(|| 42),
                    )
                    .await
                    .is_err()
                );
                match mode {
                    0 => {
                        rows.close();
                        assert_eq!(rows.next().await.unwrap(), None);
                    }
                    1 => drop(rows),
                    _ => {
                        let alias = rows.clone();
                        alias.close();
                        assert_eq!(rows.next().await.unwrap(), None);
                    }
                }
                assert_eq!(scope.blocking(|| 42).await.unwrap(), 42);
                let found = context
                    .on(0, scope)
                    .await
                    .unwrap()
                    .query(scope, "SELECT 42", vec![], 1, |row| row.get::<_, i64>(0))
                    .await
                    .unwrap();
                assert_eq!(found, [42]);
            }
            let rows = context
                .stream(0, scope, Sql::same("SELECT 'bad'"), vec![], |row| {
                    row.get::<_, i64>(0)
                })
                .await
                .unwrap();
            let StreamError::Database(error) = rows.next().await.unwrap_err() else {
                panic!("database error must retain its category")
            };
            assert_eq!(error.kind, Kind::InvalidData);
            assert_eq!(rows.next().await.unwrap(), None);
            assert!(context.on(0, scope).await.is_ok());
            drop(context);
            registry.close().await;
        })
        .unwrap();
}

#[test]
fn cancelled_stream_owner_is_drained_and_pool_remains_usable() {
    let directory = Directory::new();
    let runtime = Runtime::new(Limits {
        threads: 2,
        workers: 2,
        tasks: 16,
    })
    .unwrap();
    runtime
        .run(async |scope| {
            let registry = registry(&directory, scope).await;
            let (ready, started) = tokio::sync::oneshot::channel();
            let owned = registry.clone();
            let task = scope
                .spawn(move |child| async move {
                    let mut context = Context::new(owned);
                    let rows = context
                        .stream(0, &child, ROWS, vec![], |row| row.get::<_, i64>(0))
                        .await
                        .unwrap();
                    assert_eq!(rows.next().await.unwrap(), Some(1));
                    ready.send(()).unwrap();
                    std::future::pending::<()>().await;
                    drop(rows);
                })
                .unwrap();
            started.await.unwrap();
            drop(task);
            let mut context = Context::new(registry.clone());
            assert!(context.on(0, scope).await.is_ok());
            drop(context);
            registry.close().await;
        })
        .unwrap();
}