NoahKusaba commented on code in PR #3277:
URL: https://github.com/apache/iceberg-rust/pull/3277#discussion_r4367400067
##########
crates/iceberg/src/runtime/mod.rs:
##########
@@ -48,6 +55,29 @@ impl<T: Send + 'static> Future for JoinHandle<T> {
}
}
+/// A [`JoinHandle`] that aborts its task when dropped.
+///
+/// Dropping a plain `JoinHandle` detaches the task, which then runs on with
+/// nobody waiting for its result. Use this instead wherever the result is
+/// private to the caller, so that dropping the caller stops the work.
+pub(crate) struct AbortOnDropHandle<T>(JoinHandle<T>);
+
+impl<T> Drop for AbortOnDropHandle<T> {
+ fn drop(&mut self) {
+ self.0.0.abort();
Review Comment:
nit: `self.0.0.abort()` takes a moment to read. Could `JoinHandle` get a
`pub(crate) fn abort(&self)`, so this becomes `self.0.abort()`?
##########
crates/iceberg/src/arrow/reader/pipeline.rs:
##########
@@ -3795,4 +3811,220 @@ mod tests {
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 3);
}
+
+ /// A data file must be opened even when nothing polls the scan. The stream
+ /// stops polling the opens once it holds `concurrency_limit_data_files` of
+ /// them, and on a shared HTTP/2 connection an open stopped mid-request
holds
+ /// part of the connection window, which stalls every other read (#3276).
+ mod open_progress {
+ use std::ops::Range;
+
+ use async_trait::async_trait;
+ use bytes::Bytes;
+ use futures::StreamExt;
+ use futures::stream::BoxStream;
+ use serde::{Deserialize, Serialize};
+
+ use super::*;
+ use crate::Result;
+ use crate::arrow::scan_metrics::ScanMetrics;
+ use crate::io::{
+ FileIOBuilder, FileMetadata, FileRead, FileWrite, InputFile,
MemoryStorage, OutputFile,
+ Storage, StorageConfig, StorageFactory,
+ };
+ use crate::scan::ArrowRecordBatchStream;
+
+ const CONCURRENCY: usize = 2;
+ /// Smaller than a footer, so one open takes more than one request.
+ const METADATA_SIZE_HINT: usize = 64;
+ /// Scheduler turns, far more than a request needs to finish.
+ const YIELDS: usize = 100;
+
+ /// In-memory files whose reads take a scheduler turn to answer.
+ #[derive(Debug, Clone, Default, Serialize, Deserialize)]
+ struct CountedStorage {
Review Comment:
Since the counting moved to `ScanMetrics`, `CountedStorage` and
`CountedRead` don't count anything. All they do now is defer each read by one
scheduler turn. `DeferredStorage` / `DeferredRead` would describe them better.
More generally, most of this module is the `Storage` and `StorageFactory`
delegation plus typetag and serde. With a per-read `SpawnedFileRead`, the test
could target the wrapper directly in about 30 lines:
1. Wrap a `FileRead` that yields once and sets a flag when it finishes.
2. `poll!` the `SpawnedFileRead::read` future once and never again. Yield a
few times and assert the flag is set.
3. Do the same, but drop the future after the first poll, and assert the
flag stays unset.
That keeps the same coverage without a storage, a factory or a scan.
##########
crates/iceberg/src/runtime/mod.rs:
##########
@@ -48,6 +55,29 @@ impl<T: Send + 'static> Future for JoinHandle<T> {
}
}
+/// A [`JoinHandle`] that aborts its task when dropped.
Review Comment:
nit: Maybe point to `tokio_util::task::AbortOnDropHandle` (and DataFusion's
`SpawnedTask`) as the model here. That explains why this exists without pulling
in `tokio-util`.
##########
crates/iceberg/src/arrow/reader/pipeline.rs:
##########
@@ -3795,4 +3811,220 @@ mod tests {
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 3);
}
+
+ /// A data file must be opened even when nothing polls the scan. The stream
+ /// stops polling the opens once it holds `concurrency_limit_data_files` of
+ /// them, and on a shared HTTP/2 connection an open stopped mid-request
holds
+ /// part of the connection window, which stalls every other read (#3276).
+ mod open_progress {
+ use std::ops::Range;
+
+ use async_trait::async_trait;
+ use bytes::Bytes;
+ use futures::StreamExt;
+ use futures::stream::BoxStream;
+ use serde::{Deserialize, Serialize};
+
+ use super::*;
+ use crate::Result;
+ use crate::arrow::scan_metrics::ScanMetrics;
+ use crate::io::{
+ FileIOBuilder, FileMetadata, FileRead, FileWrite, InputFile,
MemoryStorage, OutputFile,
+ Storage, StorageConfig, StorageFactory,
+ };
+ use crate::scan::ArrowRecordBatchStream;
+
+ const CONCURRENCY: usize = 2;
+ /// Smaller than a footer, so one open takes more than one request.
+ const METADATA_SIZE_HINT: usize = 64;
+ /// Scheduler turns, far more than a request needs to finish.
+ const YIELDS: usize = 100;
Review Comment:
nit: A note that `YIELDS` assumes a current-thread runtime would help. Both
tests rely on `#[tokio::test]` being current-thread with `Runtime::current()`.
Under `flavor = "multi_thread"`, 100 yields is no longer a fixed bound and the
tests would depend on timing.
##########
crates/iceberg/src/arrow/reader/pipeline.rs:
##########
@@ -3795,4 +3811,220 @@ mod tests {
let total_rows: usize = batches.iter().map(|b| b.num_rows()).sum();
assert_eq!(total_rows, 3);
}
+
+ /// A data file must be opened even when nothing polls the scan. The stream
+ /// stops polling the opens once it holds `concurrency_limit_data_files` of
+ /// them, and on a shared HTTP/2 connection an open stopped mid-request
holds
+ /// part of the connection window, which stalls every other read (#3276).
+ mod open_progress {
+ use std::ops::Range;
+
+ use async_trait::async_trait;
+ use bytes::Bytes;
+ use futures::StreamExt;
+ use futures::stream::BoxStream;
+ use serde::{Deserialize, Serialize};
+
+ use super::*;
+ use crate::Result;
+ use crate::arrow::scan_metrics::ScanMetrics;
+ use crate::io::{
+ FileIOBuilder, FileMetadata, FileRead, FileWrite, InputFile,
MemoryStorage, OutputFile,
+ Storage, StorageConfig, StorageFactory,
+ };
+ use crate::scan::ArrowRecordBatchStream;
+
+ const CONCURRENCY: usize = 2;
+ /// Smaller than a footer, so one open takes more than one request.
+ const METADATA_SIZE_HINT: usize = 64;
+ /// Scheduler turns, far more than a request needs to finish.
+ const YIELDS: usize = 100;
+
+ /// In-memory files whose reads take a scheduler turn to answer.
+ #[derive(Debug, Clone, Default, Serialize, Deserialize)]
+ struct CountedStorage {
+ #[serde(skip)]
+ files: MemoryStorage,
+ }
+
+ struct CountedRead {
+ contents: Bytes,
+ }
+
+ #[async_trait]
+ impl FileRead for CountedRead {
+ async fn read(&self, range: Range<u64>) -> Result<Bytes> {
+ // A request in flight when the scan stops being polled
finishes
+ // only if something else drives it.
+ tokio::task::yield_now().await;
+ Ok(self
+ .contents
+ .slice(range.start as usize..range.end as usize))
+ }
+ }
+
+ #[async_trait]
+ #[typetag::serde]
+ impl Storage for CountedStorage {
+ async fn reader(&self, path: &str) -> Result<Box<dyn FileRead>> {
+ Ok(Box::new(CountedRead {
+ contents: self.files.read(path).await?,
+ }))
+ }
+
+ fn new_input(&self, path: &str) -> Result<InputFile> {
+ Ok(InputFile::new(Arc::new(self.clone()), path.to_string()))
+ }
+
+ async fn exists(&self, path: &str) -> Result<bool> {
+ self.files.exists(path).await
+ }
+
+ async fn metadata(&self, path: &str) -> Result<FileMetadata> {
+ self.files.metadata(path).await
+ }
+
+ async fn read(&self, path: &str) -> Result<Bytes> {
+ self.files.read(path).await
+ }
+
+ async fn write(&self, path: &str, bs: Bytes) -> Result<()> {
+ self.files.write(path, bs).await
+ }
+
+ async fn writer(&self, path: &str) -> Result<Box<dyn FileWrite>> {
+ self.files.writer(path).await
+ }
+
+ async fn delete(&self, path: &str) -> Result<()> {
+ self.files.delete(path).await
+ }
+
+ async fn delete_prefix(&self, path: &str) -> Result<()> {
+ self.files.delete_prefix(path).await
+ }
+
+ async fn delete_stream(&self, paths: BoxStream<'static, String>)
-> Result<()> {
+ self.files.delete_stream(paths).await
+ }
+
+ fn new_output(&self, path: &str) -> Result<OutputFile> {
+ Ok(OutputFile::new(Arc::new(self.clone()), path.to_string()))
+ }
+ }
+
+ #[derive(Debug, Default, Serialize, Deserialize)]
+ struct CountedStorageFactory {
+ #[serde(skip)]
+ storage: CountedStorage,
+ }
+
+ #[typetag::serde]
+ impl StorageFactory for CountedStorageFactory {
+ fn build(&self, _config: &StorageConfig) -> Result<Arc<dyn
Storage>> {
+ Ok(Arc::new(self.storage.clone()))
+ }
+ }
+
+ /// A scan of `CONCURRENCY` identical Parquet files, and the metrics
that
+ /// count the bytes it has read.
+ async fn counted_scan() -> (ScanMetrics, ArrowRecordBatchStream) {
+ let schema = Arc::new(
+ Schema::builder()
+ .with_schema_id(1)
+ .with_fields(vec![
+ NestedField::required(1, "id",
Type::Primitive(PrimitiveType::Int)).into(),
+ ])
+ .build()
+ .unwrap(),
+ );
+ let tmp_dir = TempDir::new().unwrap();
+ let on_disk = write_plain_parquet(
+ tmp_dir.path().to_str().unwrap(),
+ "data.parquet",
+ vec![],
+ vec![],
+ );
+ let contents = Bytes::from(std::fs::read(on_disk).unwrap());
+
+ let storage = CountedStorage::default();
+ let mut tasks = Vec::with_capacity(CONCURRENCY);
+ for file_num in 0..CONCURRENCY {
+ let path = format!("memory:/data/file_{file_num}.parquet");
+ storage.write(&path, contents.clone()).await.unwrap();
+ tasks.push(
+ FileScanTask::builder()
+ .with_data_file_path(path)
+ .with_file_size_in_bytes(contents.len() as u64)
+ .with_start(0)
+ .with_length(0)
+ .with_data_file_format(DataFileFormat::Parquet)
+ .with_schema(Arc::clone(&schema))
+ .with_project_field_ids(vec![1])
+ .with_case_sensitive(false)
+ .build()
+ .unwrap(),
+ );
+ }
+
+ let file_io = FileIOBuilder::new(Arc::new(CountedStorageFactory {
storage })).build();
+ let reader = ArrowReaderBuilder::new(file_io, Runtime::current())
+ .with_data_file_concurrency_limit(CONCURRENCY)
+ .with_metadata_size_hint(METADATA_SIZE_HINT)
+ .build();
+ let tasks =
+ Box::pin(futures::stream::iter(tasks.into_iter().map(Ok))) as
FileScanTaskStream;
+ let result = reader.read(tasks).unwrap();
+ (result.metrics().clone(), result.stream())
+ }
+
+ /// Gives the scheduler turns until `done` holds, or `YIELDS` pass.
+ async fn spin(done: impl Fn() -> bool) {
+ for _ in 0..YIELDS {
+ if done() {
+ return;
+ }
+ tokio::task::yield_now().await;
+ }
+ }
+
+ /// An open must carry on through its requests although the scan is
never
+ /// polled again.
+ #[tokio::test]
+ async fn test_open_proceeds_without_the_scan_being_polled() {
+ let (metrics, scan) = counted_scan().await;
+ let mut scan = Box::pin(scan);
+
+ // One poll starts an open per concurrency slot. Nothing polls the
+ // scan after this, so only a spawned open can read any further.
+ let _ = futures::poll!(scan.next());
+ let on_the_first_poll = metrics.bytes_read();
+ spin(|| metrics.bytes_read() > on_the_first_poll).await;
+ assert!(
+ metrics.bytes_read() > on_the_first_poll,
Review Comment:
`bytes_read() > on_the_first_poll` passes as soon as one more request goes
out. The inputs are fixed (`CONCURRENCY` identical files and a fixed metadata
hint), so could this wait until the opens finish and `assert_eq!` the exact
total? The drop test already does this.
##########
crates/iceberg/src/runtime/mod.rs:
##########
@@ -48,6 +55,29 @@ impl<T: Send + 'static> Future for JoinHandle<T> {
}
}
+/// A [`JoinHandle`] that aborts its task when dropped.
+///
+/// Dropping a plain `JoinHandle` detaches the task, which then runs on with
+/// nobody waiting for its result. Use this instead wherever the result is
+/// private to the caller, so that dropping the caller stops the work.
+pub(crate) struct AbortOnDropHandle<T>(JoinHandle<T>);
+
+impl<T> Drop for AbortOnDropHandle<T> {
+ fn drop(&mut self) {
+ self.0.0.abort();
+ }
+}
+
+impl<T> Unpin for AbortOnDropHandle<T> {}
Review Comment:
nit: This impl isn't needed. `JoinHandle<T>` is `Unpin` for every `T`, so
`AbortOnDropHandle<T>` is `Unpin` automatically.
##########
crates/iceberg/src/arrow/reader/pipeline.rs:
##########
@@ -144,16 +147,29 @@ impl FileScanTaskReader {
.delete_file_loader
.load_deletes(task.deletes(), task.schema_ref());
- // Open the Parquet file once, loading its metadata
- let (parquet_file_reader, arrow_metadata) =
ArrowReader::open_parquet_file(
- task.data_file_path(),
- &self.file_io,
- task.file_size_in_bytes(),
- parquet_read_options,
- self.scan_metrics.bytes_read_counter(),
- task.key_metadata(),
- )
- .await?;
+ // Spawned so the open finishes even when the caller stops polling
(#3276).
Review Comment:
nit: The caller is still polling. What stops is the flatten stage polling
the buffered opens once it holds `N` streams. Maybe:
```rust
// Spawned because `try_flatten_unordered` stops polling the buffered opens
once it
// holds `concurrency_limit_data_files` streams, leaving requests half read
(#3276).
```
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
