BufWriter panic #810

thread 'writer::test::test::test_async_arrow_writer_close_after_write_error_panics' (198483) panicked at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/object_store-0.13.2/src/buffered.rs:407:71:
`async fn` resumed after completion
stack backtrace:
   0: __rustc::rust_begin_unwind
             at /rustc/ac68faa20c58cbccd01ee7208bf3b6e93a7d7f96/library/std/src/panicking.rs:689:5
   1: core::panicking::panic_fmt
             at /rustc/ac68faa20c58cbccd01ee7208bf3b6e93a7d7f96/library/core/src/panicking.rs:80:14
   2: core::panicking::panic_const::panic_const_async_fn_resumed
             at /rustc/ac68faa20c58cbccd01ee7208bf3b6e93a7d7f96/library/core/src/panicking.rs:175:17
   3: <object_store::buffered::BufWriter as tokio::io::async_write::AsyncWrite>::poll_write::{{closure}}
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/object_store-0.13.2/src/buffered.rs:407:71
   4: <core::pin::Pin<P> as core::future::future::Future>::poll
             at /home/xxx/.rust/toolchains/1.96.0-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/future/future.rs:133:9
   5: futures_util::future::future::FutureExt::poll_unpin
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/futures-util-0.3.32/src/future/future/mod.rs:556:24
   6: <object_store::buffered::BufWriter as tokio::io::async_write::AsyncWrite>::poll_write
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/object_store-0.13.2/src/buffered.rs:394:65
   7: <&mut T as tokio::io::async_write::AsyncWrite>::poll_write
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/tokio-1.52.3/src/io/async_write.rs:186:35
   8: <tokio::io::util::write_all::WriteAll<W> as core::future::future::Future>::poll
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/tokio-1.52.3/src/io/util/write_all.rs:43:54
   9: <T as parquet::arrow::async_writer::AsyncFileWriter>::write::{{closure}}
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parquet-58.3.0/src/arrow/async_writer/mod.rs:109:33
  10: <core::pin::Pin<P> as core::future::future::Future>::poll
             at /home/xxx/.rust/toolchains/1.96.0-x86_64-unknown-linux-gnu/lib/rustlib/src/rust/library/core/src/future/future.rs:133:9
  11: parquet::arrow::async_writer::AsyncArrowWriter<W>::do_write::{{closure}}
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parquet-58.3.0/src/arrow/async_writer/mod.rs:288:14
  12: parquet::arrow::async_writer::AsyncArrowWriter<W>::write::{{closure}}
             at /home/xxx/.cargo/registry/src/index.crates.io-1949cf8c6b5b557f/parquet-58.3.0/src/arrow/async_writer/mod.rs:223:29

#[cfg(test)]
mod test {
    use std::sync::Arc;

    use arrow_array::{Int64Array, RecordBatch};
    use object_store::Error;
    use object_store::buffered::BufWriter;
    use object_store::local::LocalFileSystem;
    use object_store::path::Path;
    use object_store::{
        CopyOptions, GetOptions, GetResult, ListResult, MultipartUpload, ObjectMeta, ObjectStore,
        PutMultipartOptions, PutOptions, PutPayload, PutResult, Result,
    };
    use arrow_schema::{DataType, Field, Schema};
    use bytes::Bytes;
    use futures_util::stream::BoxStream;
    use parquet::{
        arrow::AsyncArrowWriter, basic::Compression, file::properties::WriterProperties,
    };
    use tokio::io::AsyncWriteExt;

    #[derive(Debug)]
    struct FailMultipartInitStore {
        inner: Arc<LocalFileSystem>,
    }

    impl FailMultipartInitStore {
        fn new(inner: Arc<LocalFileSystem>) -> Self {
            Self { inner }
        }
    }

    impl std::fmt::Display for FailMultipartInitStore {
        fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
            write!(f, "FailMultipartInitStore")
        }
    }

    #[async_trait::async_trait]
    impl ObjectStore for FailMultipartInitStore {
        async fn put_opts(
            &self,
            location: &Path,
            payload: PutPayload,
            opts: PutOptions,
        ) -> Result<PutResult> {
            self.inner.put_opts(location, payload, opts).await
        }

        async fn put_multipart_opts(
            &self,
            _location: &Path,
            _opts: PutMultipartOptions,
        ) -> Result<Box<dyn MultipartUpload>> {
            Err(Error::Generic {
                store: "FailMultipartInitStore",
                source: Box::new(std::io::Error::other("StorageFull")),
            })
        }

        async fn get_opts(&self, location: &Path, options: GetOptions) -> Result<GetResult> {
            self.inner.get_opts(location, options).await
        }

        fn delete_stream(
            &self,
            locations: BoxStream<'static, Result<Path>>,
        ) -> BoxStream<'static, Result<Path>> {
            self.inner.delete_stream(locations)
        }

        fn list(&self, prefix: Option<&Path>) -> BoxStream<'static, Result<ObjectMeta>> {
            self.inner.list(prefix)
        }

        async fn list_with_delimiter(&self, prefix: Option<&Path>) -> Result<ListResult> {
            self.inner.list_with_delimiter(prefix).await
        }

        async fn copy_opts(&self, from: &Path, to: &Path, options: CopyOptions) -> Result<()> {
            self.inner.copy_opts(from, to, options).await
        }
    }

    #[tokio::test]
    #[should_panic(expected = "`async fn` resumed after completion")]
    async fn test_buf_writer_shutdown_after_write_error_panics() {
        let dir = tempfile::tempdir().unwrap();
        let dir = dir.path().to_str().unwrap();
        let inner = Arc::new(LocalFileSystem::new_with_prefix(dir).unwrap());
        let store = Arc::new(FailMultipartInitStore::new(inner));
        let path = Path::from("panic-repro.parquet");
        let mut writer = BufWriter::with_capacity(store, path, 1);

        let err = writer
            .write_all_buf(&mut Bytes::from_static(b"trigger-prepare-error"))
            .await;
        assert!(
            err.is_err(),
            "expected write_all_buf to fail before shutdown"
        );

        let _ = writer.shutdown().await;
    }

    #[tokio::test]
    #[should_panic(expected = "`async fn` resumed after completion")]
    async fn test_async_arrow_writer_close_after_write_error_panics() {
        let dir = tempfile::tempdir().unwrap();
        let dir = dir.path().to_str().unwrap();
        let inner = Arc::new(LocalFileSystem::new_with_prefix(dir).unwrap());
        let store = Arc::new(FailMultipartInitStore::new(inner));
        let path = Path::from("panic-repro-arrow-writer.parquet");
        let buf_writer = BufWriter::with_capacity(store, path, 1);

        let schema = Arc::new(Schema::new(vec![Field::new(
            "value",
            DataType::Int64,
            false,
        )]));
        let writer_props = WriterProperties::builder()
            .set_compression(Compression::UNCOMPRESSED)
            .set_max_row_group_row_count(Some(1))
            .build();
        let mut writer =
            AsyncArrowWriter::try_new(buf_writer, schema.clone(), Some(writer_props)).unwrap();

        let values = Int64Array::from_iter_values([1_i64]);
        let batch = RecordBatch::try_new(schema, vec![Arc::new(values)]).unwrap();

        for _ in 0..100_000 {
            let _ = writer.write(&batch).await;
        }

        let _ = writer.close().await;
    }
}