async: add rados_async_write_stream -> WriteSink

This helps writers who would like to stream
data into a RADOS object without managing their own
backlog of futures.

Author: jcsp

src/ceph.rs

@@ -43,6 +43,7 @@ use std::time::{Duration, SystemTime, UNIX_EPOCH};
 
 use crate::list_stream::ListStream;
 use crate::read_stream::ReadStream;
+pub use crate::write_sink::WriteSink;
 use std::pin::Pin;
 use std::sync::Arc;
 use std::task::{Context, Poll};
@@ -1076,10 +1077,47 @@ impl IoCtx {
     /// Streaming read of a RADOS object.  The `ReadStream` object implements `futures::Stream`
     /// for use with Stream-aware code like hyper's Body::wrap_stream.
     ///
-    /// This will usually issue more read ops than needed if used on a small object: for
-    /// small objects `rados_async_object_read` is more appropriate.
-    pub fn rados_async_object_read_stream(&self, object_name: &str) -> ReadStream<'_> {
-        ReadStream::new(self, object_name, None, None)
+    /// Useful for reading large objects incrementally, or anywhere you are using an interface
+    /// that expects a stream (such as proxying objects via an HTTP server).
+    ///
+    /// Efficiency: If size_hint is not specified, and this function is used on a small object, it will
+    /// issue spurious read-ahead operations beyond the object's size.
+    /// If you have an object that you know is small, prefer to use a single `rados_async_object_read`
+    /// instead of this streaming variant.
+    ///
+    /// * `buffer_size` - How much data should be read per rados read operation.  This is also
+    ///   how much data is emitted in each Item from the stream.
+    /// * `concurrency` - How many RADOS operations should be run in parallel for this stream,
+    ///   or None to use a default.
+    /// * `size_hint` - If you have prior knowledge of the object's size in bytes, pass it here to enable
+    ///   the stream to issue fewer read-ahead operations than it would by default.  This is just
+    ///   a hint, and does not bound the data returned -- if the object is smaller or larger
+    ///   than `size_hint` then the actual object size will be reflected in the stream's output.
+    pub fn rados_async_object_read_stream(
+        &self,
+        object_name: &str,
+        buffer_size: Option<usize>,
+        concurrency: Option<usize>,
+        size_hint: Option<u64>,
+    ) -> ReadStream<'_> {
+        ReadStream::new(self, object_name, buffer_size, concurrency, size_hint)
+    }
+
+    /// Streaming write of a RADOS object.  The `WriteSink` object implements `futures::Sink`.  Combine
+    /// it with other stream-aware code, or bring the SinkExt trait into scope to get methods
+    /// like send, send_all.
+    ///
+    /// Efficiency: this class does not coalesce writes, so each Item you send into it,
+    ///
+    ///
+    /// * `concurrency` - How many RADOS operations should be run in parallel for this stream,
+    ///   or None to use a default.
+    pub fn rados_async_object_write_stream(
+        &self,
+        object_name: &str,
+        concurrency: Option<usize>,
+    ) -> WriteSink<'_> {
+        WriteSink::new(self, object_name, concurrency)
     }
 
     /// Get object stats (size,SystemTime)

src/lib.rs

@@ -82,6 +82,7 @@ pub(crate) mod completion;
 pub(crate) mod list_stream;
 mod mon_command;
 pub(crate) mod read_stream;
+pub(crate) mod write_sink;
 
 pub use crate::ceph_client::CephClient;
 pub use crate::ceph_version::CephVersion;

src/read_stream.rs

@@ -36,10 +36,18 @@ pub struct ReadStream<'a> {
     // Number of concurrent RADOS read ops to issue
     concurrency: usize,
 
+    // Caller's hint as to the object size (not required to be accurate)
+    size_hint: Option<u64>,
+
     in_flight: Vec<IOSlot<'a>>,
 
+    // Counter for how many bytes we have issued reads for
     next: u64,
 
+    // Counter for how many bytes we have yielded from poll_next()
+    // (i.e. the size of the object so far)
+    yielded: u64,
+
     object_name: String,
 
     // Flag is set when we see a short read - means do not issue any more IOs,
@@ -55,13 +63,16 @@ impl<'a> ReadStream<'a> {
         object_name: &str,
         buffer_size: Option<usize>,
         concurrency: Option<usize>,
+        size_hint: Option<u64>,
     ) -> Self {
         let mut inst = Self {
             ioctx,
             buffer_size: buffer_size.unwrap_or(DEFAULT_BUFFER_SIZE),
             concurrency: concurrency.unwrap_or(DEFAULT_CONCURRENCY),
+            size_hint,
             in_flight: Vec::new(),
             next: 0,
+            yielded: 0,
             object_name: object_name.to_string(),
             done: false,
         };
@@ -80,8 +91,20 @@ enum IOSlot<'a> {
 
 impl<'a> ReadStream<'a> {
     fn maybe_issue(&mut self) {
-        // Issue reads
-        while self.in_flight.len() < self.concurrency {
+        // Issue reads if any of these are true:
+        // - Nothing is in flight
+        // - No size bound, and in flight < concurrency
+        // - A size bound, and we're within it, and in flight < concurrency
+        // - A size bound, and it has been disproved, and in flight < concurrency
+
+        while !self.done
+            && (self.in_flight.is_empty()
+                || (((self.size_hint.is_some()
+                    && (self.next < self.size_hint.unwrap()
+                        || self.yielded > self.size_hint.unwrap()))
+                    || self.size_hint.is_none())
+                    && (self.in_flight.len() < self.concurrency)))
+        {
             let read_at = self.next;
             self.next += self.buffer_size as u64;
 
@@ -163,10 +186,8 @@ impl<'a> Stream for ReadStream<'a> {
             },
         };
 
-        self.maybe_issue();
-
         // A result is ready, handle it.
-        match result {
+        let r = match result {
             Ok(length) => {
                 if (length as usize) < self.buffer_size {
                     // Cancel outstanding ops
@@ -175,9 +196,14 @@ impl<'a> Stream for ReadStream<'a> {
                     // Flag to return Ready(None) on next call to poll.
                     self.done = true;
                 }
+                self.yielded += buffer.len() as u64;
                 Poll::Ready(Some(Ok(buffer)))
             }
             Err(e) => Poll::Ready(Some(Err(e))),
-        }
+        };
+
+        self.maybe_issue();
+
+        r
     }
 }

src/write_sink.rs

@@ -0,0 +1,123 @@
+use futures::{FutureExt, Sink, Stream};
+use std::future::Future;
+use std::pin::Pin;
+use std::task::{Context, Poll};
+
+use crate::ceph::IoCtx;
+use crate::completion::with_completion;
+use crate::error::{RadosError, RadosResult};
+use crate::rados::rados_aio_write;
+use futures::stream::FuturesUnordered;
+use std::ffi::CString;
+use std::os::raw::c_char;
+
+const DEFAULT_CONCURRENCY: usize = 2;
+
+pub struct WriteSink<'a> {
+    ioctx: &'a IoCtx,
+    in_flight: Pin<Box<FuturesUnordered<Pin<Box<dyn Future<Output = RadosResult<u32>> + 'a>>>>>,
+    object_name: String,
+
+    // Offset into object where the next write will land
+    next: u64,
+
+    // How many RADOS ops in flight at same time?
+    concurrency: usize,
+}
+
+unsafe impl Send for WriteSink<'_> {}
+
+impl<'a> WriteSink<'a> {
+    pub fn new(ioctx: &'a IoCtx, object_name: &str, concurrency: Option<usize>) -> Self {
+        let concurrency = concurrency.unwrap_or(DEFAULT_CONCURRENCY);
+        assert!(concurrency > 0);
+
+        Self {
+            ioctx,
+            in_flight: Box::pin(FuturesUnordered::new()),
+            object_name: object_name.to_string(),
+            next: 0,
+            concurrency,
+        }
+    }
+
+    fn trim_in_flight(
+        mut self: Pin<&mut Self>,
+        cx: &mut Context<'_>,
+        target_len: usize,
+    ) -> Poll<Result<(), <Self as Sink<Vec<u8>>>::Error>> {
+        while self.in_flight.len() > target_len {
+            match self.in_flight.as_mut().poll_next(cx) {
+                Poll::Pending => return Poll::Pending,
+                Poll::Ready(None) => {
+                    // (because we check for in_flight size first)
+                    unreachable!()
+                }
+                Poll::Ready(Some(result)) => match result {
+                    Err(e) => return Poll::Ready(Err(e)),
+                    Ok(sz) => {
+                        debug!("trim_in_flight: IO completed with r={}", sz);
+                    }
+                },
+            };
+        }
+
+        // Nothing left in flight, we're done
+        Poll::Ready(Ok(()))
+    }
+}
+
+impl<'a> Sink<Vec<u8>> for WriteSink<'a> {
+    type Error = RadosError;
+
+    fn poll_ready(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        // If we have fewer than 1 slots available, this will try to wait on some outstanding futures
+        let target = self.as_ref().concurrency - 1;
+        if self.in_flight.len() > target {
+            self.trim_in_flight(cx, target)
+        } else {
+            Poll::Ready(Ok(()))
+        }
+    }
+
+    fn start_send(mut self: Pin<&mut Self>, item: Vec<u8>) -> Result<(), Self::Error> {
+        let ioctx = self.ioctx;
+        let obj_name_str = CString::new(self.object_name.clone()).expect("CString error");
+        let write_at = self.next;
+        self.next += item.len() as u64;
+
+        let mut fut = Box::pin(async move {
+            let c = with_completion(ioctx, |c| unsafe {
+                rados_aio_write(
+                    ioctx.ioctx,
+                    obj_name_str.as_ptr(),
+                    c,
+                    item.as_ptr() as *mut c_char,
+                    item.len(),
+                    write_at,
+                )
+            })?;
+
+            c.await
+        });
+
+        // Kick the async{} future to get the RADOS op sent
+        match fut.as_mut().now_or_never() {
+            Some(Ok(_)) => Ok(()),
+            Some(Err(e)) => return Err(e),
+            None => {
+                self.in_flight.push(fut);
+                Ok(())
+            }
+        }
+    }
+
+    fn poll_flush(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        self.trim_in_flight(cx, 0)
+    }
+
+    fn poll_close(self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll<Result<(), Self::Error>> {
+        // There is no special work to be done on close
+        self.poll_flush(cx)
+    }
+}