Object Store API improvements

[fsdvh]

Describe the bug

So recently we started seeing two issues:

• multiple shutdowns of the same writer will cause an issue, related to [this](

  arrow-rs/object_store/src/buffered.rs

  Line 220 in 50e9e49

  pub struct BufWriter {

• partition size issue, related to this

Multiple Shutdown

First of all, I want to note that multiple shutdown calls to the same writer are the issue by itself, but I think we can make the situation better with minimum effort.

Here is the code:

             BufWriterState::Write(x) => {
                    let upload = x.take().ok_or_else(|| {
                        std::io::Error::new(
                            ErrorKind::InvalidInput,
                            "Cannot shutdown a writer that has already been shut down",
                        )
                    })?;
                    self.state = BufWriterState::Flush(
                        async move {
                            upload.finish().await?;
                            Ok(())
                        }
                        .boxed(),
                    )
                }

I think we can change it to something more friendly like this:

             BufWriterState::Write(x) => {
                    if let Some(upload) = x.take() {
                        self.state = BufWriterState::Flush(
                            async move { upload.finish().await.map(|_| ()) }.boxed(),
                        )
                    } else {
                        return Poll::Ready(Ok(()));
                    }
                }

This way on a second shutdown call we just immediately return Ok(())

Upload part size issue

Something leftover during the shutdown, complete before the previous upload, in this case, we're getting:

Your proposed upload is smaller than the minimum allowed size

To mitigate this issue we probably should wait for all previous part uploads to complete and then upload the final part which may be smaller than the minimum size of the last one.

Here is the original code I propose to change:

 pub async fn finish(mut self) -> Result<PutResult> {
        if !self.buffer.is_empty() {
            let part = std::mem::take(&mut self.buffer);
            self.put_part(part.into())
        }

        self.wait_for_capacity(0).await?;

        match self.upload.complete().await {
            Err(e) => {
                self.tasks.shutdown().await;
                self.upload.abort().await?;
                Err(e)
            }
            Ok(result) => Ok(result),
        }
    }

by injecting self.wait_for_capacity(0).await?; before actually putting the last chunk we can mitigate this issue.

 pub async fn finish(mut self) -> Result<PutResult> {
        if !self.buffer.is_empty() {
            self.wait_for_capacity(0).await?;  // here

            let part = std::mem::take(&mut self.buffer);
            self.put_part(part.into())
        }

        self.wait_for_capacity(0).await?;

        match self.upload.complete().await {
            Err(e) => {
                self.tasks.shutdown().await;
                self.upload.abort().await?;
                Err(e)
            }
            Ok(result) => Ok(result),
        }
    }

This way we wait for all ongoing uploads before submitting the last part

[fsdvh]

I think it's more an enhancement

[tustvold]

What is the use-case for multiple shutdowns, this feels like a bug in the calling code?

A similar argument could be made for calling shutdown before the parts have been uploaded, this I think implies the caller is not waiting on the returned futures?

I'm actually confused by the second one, could you provide more context on how you encountered this error, including which store implementation you are using. Parts are numbered and so it shouldn't matter if they complete out of order.

[fsdvh]

Yes, calling multiple shutdowns it's a bug on the caller side for sure. Changes proposed by me should help mitigate this issue a bit by making shutdown calls relaxed

For the second one, we're using an s3 object store and for some reason, it (s3) doesn't take into account the part number during the multipart upload and return error mentioned above. We managed to solve this issue by calling self.wait_for_capacity(0).await? before calling finish()

[tustvold]

Changes proposed by me should help mitigate this issue a bit by making shutdown calls relaxed

Your proposed change is racy, the first close will wait on the upload to complete, with subsequent calls "completing" instantly.

For the second one, we're using an s3 object store and for some reason, it (s3) doesn't take into account the part number during the multipart upload and return error mentioned above. We managed to solve this issue by calling self.wait_for_capacity(0).await? before calling finish()

Perhaps you could write a simple reproducer for this, I'm not saying S3 doesn't do this but I want to be sure we've correctly identified the issue. Your proposed fix will serialize a round-trip which is unfortunate when many stores are not exactly low-latency.

[fsdvh]

I will try to provide an example of the second issue, but meanwhile, I thought maybe we can change the flush method of BufWriter:

fn poll_flush(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Error>> {
        loop {
            return match &mut self.state {
                 BufWriterState::Write(write) => {
                    if let Some(write) = write {
                        write.poll_for_capacity(cx, 0).map_err(|e| e.into())
                    } else {
                        panic!("Already shut down")
                    }
                },
                BufWriterState::Buffer(_, _) => Poll::Ready(Ok(())),
                BufWriterState::Flush(_) => panic!("Already shut down"),
                BufWriterState::Prepare(f) => {
                    self.state = BufWriterState::Write(ready!(f.poll_unpin(cx)?).into());
                    continue;
                }
            };
        }
    }

By actually waiting for all downloads to complete we can give a user and ability to use flush() + shutdown(), wdyt?

[tustvold]

What would be the benefit of this? It wouldn't be able to guarantee that there isn't any data still in flight as it can't upload the final data until shutdown is called. It is a valid point that we're taking a somewhat dubious interpretation of the AsyncWrite trait, but its a necessary evil. We could probably add further documentation to discourage the use of AsyncWrite

[fsdvh]

yes, but flush will ensure that ongoing writes are flushed which is generally aligned with "flush". And after that, you're okay to finalize your upload in case of shutdown

[tustvold]

will ensure that ongoing writes are flushed which is generally aligned with "flush"

Only those it has actually started writing, I think the proposed behaviour is more confusing. We should just make shutdown do the right thing