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 { 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(); }