RUST-1588 Implement IterateOnce unified test runner operation (#936)

Author: isabelatkinson

src/cursor.rs

@@ -208,6 +208,16 @@ impl<T> Cursor<T> {
         self.wrapped_cursor.as_mut().unwrap().advance().await
     }
 
+    #[cfg(test)]
+    pub(crate) async fn try_advance(&mut self) -> Result<()> {
+        self.wrapped_cursor
+            .as_mut()
+            .unwrap()
+            .try_advance()
+            .await
+            .map(|_| ())
+    }
+
     /// Returns a reference to the current result in the cursor.
     ///
     /// # Panics

src/cursor/common.rs

@@ -24,6 +24,16 @@ use crate::{
     Namespace,
 };
 
+/// The result of one attempt to advance a cursor.
+pub(super) enum AdvanceResult {
+    /// The cursor was successfully advanced and the buffer has at least one item.
+    Advanced,
+    /// The cursor does not have any more items and will not return any more in the future.
+    Exhausted,
+    /// The cursor does not currently have any items, but future calls to getMore may yield more.
+    Waiting,
+}
+
 /// An internal cursor that can be used in a variety of contexts depending on its `GetMoreProvider`.
 #[derive(Derivative)]
 #[derivative(Debug)]
@@ -95,33 +105,44 @@ where
         self.state.as_ref().unwrap()
     }
 
-    /// Advance the cursor forward to the next document.
-    /// If there are no documents cached locally, perform getMores until
-    /// the cursor is exhausted or a result/error has been received.
+    /// Attempt to advance the cursor forward to the next item. If there are no items cached
+    /// locally, perform getMores until the cursor is exhausted or the buffer has been refilled.
+    /// Return whether or not the cursor has been advanced.
     pub(super) async fn advance(&mut self) -> Result<bool> {
         loop {
-            self.state_mut().buffer.advance();
-
-            if !self.state().buffer.is_empty() {
-                break;
+            match self.try_advance().await? {
+                AdvanceResult::Advanced => return Ok(true),
+                AdvanceResult::Exhausted => return Ok(false),
+                AdvanceResult::Waiting => continue,
             }
+        }
+    }
 
-            // if moving the offset puts us at the end of the buffer, perform another
-            // getMore if the cursor is still alive.
+    /// Attempt to advance the cursor forward to the next item. If there are no items cached
+    /// locally, perform a single getMore to attempt to retrieve more.
+    pub(super) async fn try_advance(&mut self) -> Result<AdvanceResult> {
+        if self.state_mut().buffer.advance() {
+            return Ok(AdvanceResult::Advanced);
+        } else if self.is_exhausted() {
+            return Ok(AdvanceResult::Exhausted);
+        }
 
-            if self.state().exhausted {
-                return Ok(false);
-            }
+        // If the buffer is empty but the cursor is not exhausted, perform a getMore.
+        let client = self.client.clone();
+        let spec = self.info.clone();
+        let pin = self.state().pinned_connection.replicate();
 
-            let client = self.client.clone();
-            let spec = self.info.clone();
-            let pin = self.state().pinned_connection.replicate();
+        let result = self.provider.execute(spec, client, pin).await;
+        self.handle_get_more_result(result)?;
 
-            let result = self.provider.execute(spec, client, pin).await;
-            self.handle_get_more_result(result)?;
+        if self.is_exhausted() {
+            Ok(AdvanceResult::Exhausted)
+        } else {
+            match self.state_mut().buffer.advance() {
+                true => Ok(AdvanceResult::Advanced),
+                false => Ok(AdvanceResult::Waiting),
+            }
         }
-
-        Ok(true)
     }
 
     pub(super) fn take_state(&mut self) -> CursorState {
@@ -494,21 +515,27 @@ impl CursorBuffer {
         self.docs.is_empty()
     }
 
+    /// Removes and returns the document in the front of the buffer.
     pub(crate) fn next(&mut self) -> Option<RawDocumentBuf> {
         self.fresh = false;
         self.docs.pop_front()
     }
 
-    pub(crate) fn advance(&mut self) {
-        // if at the front of the buffer, don't move forward as the first document
-        // hasn't been consumed yet.
+    /// Advances the buffer to the next document. Returns whether there are any documents remaining
+    /// in the buffer after advancing.
+    pub(crate) fn advance(&mut self) -> bool {
+        // If at the front of the buffer, don't move forward as the first document hasn't been
+        // consumed yet.
         if self.fresh {
             self.fresh = false;
-            return;
+        } else {
+            self.docs.pop_front();
         }
-        self.next();
+        !self.is_empty()
     }
 
+    /// Returns the item at the front of the buffer, if there is one. This method does not change
+    /// the state of the buffer.
     pub(crate) fn current(&self) -> Option<&RawDocument> {
         self.docs.front().map(|d| d.as_ref())
     }
@@ -519,3 +546,24 @@ impl AsRef<VecDeque<RawDocumentBuf>> for CursorBuffer {
         &self.docs
     }
 }
+
+#[test]
+fn test_buffer() {
+    use bson::rawdoc;
+
+    let queue: VecDeque<RawDocumentBuf> =
+        [rawdoc! { "x": 1 }, rawdoc! { "x": 2 }, rawdoc! { "x": 3 }].into();
+    let mut buffer = CursorBuffer::new(queue);
+
+    assert!(buffer.advance());
+    assert_eq!(buffer.current(), Some(rawdoc! { "x": 1 }.as_ref()));
+
+    assert!(buffer.advance());
+    assert_eq!(buffer.current(), Some(rawdoc! { "x": 2 }.as_ref()));
+
+    assert!(buffer.advance());
+    assert_eq!(buffer.current(), Some(rawdoc! { "x": 3 }.as_ref()));
+
+    assert!(!buffer.advance());
+    assert_eq!(buffer.current(), None);
+}

src/cursor/session.rs

@@ -245,6 +245,15 @@ impl<T> SessionCursor<T> {
         self.make_stream(session).generic_cursor.advance().await
     }
 
+    #[cfg(test)]
+    pub(crate) async fn try_advance(&mut self, session: &mut ClientSession) -> Result<()> {
+        self.make_stream(session)
+            .generic_cursor
+            .try_advance()
+            .await
+            .map(|_| ())
+    }
+
     /// Returns a reference to the current result in the cursor.
     ///
     /// # Panics

src/test/spec/run_command.rs

@@ -6,10 +6,8 @@ async fn run_unified() {
     let _guard = LOCK.run_exclusively().await;
     run_unified_tests(&["run-command", "unified"])
         .skip_tests(&[
-            // TODO re: RUST-1649: fix withTransaction for new test runner
+            // TODO RUST-1649: unskip when withTransaction is implemented
             "attaches transaction fields to given command",
-            // TODO fix in a follow up PR because need to add IterateOnce
-            "does not close the cursor when receiving an empty batch",
         ])
         .await;
 }

src/test/spec/unified_runner/operation.rs

@@ -382,6 +382,7 @@ impl<'de> Deserialize<'de> for Operation {
             "getKeys" => deserialize_op::<GetKeys>(definition.arguments),
             #[cfg(feature = "in-use-encryption-unstable")]
             "removeKeyAltName" => deserialize_op::<RemoveKeyAltName>(definition.arguments),
+            "iterateOnce" => deserialize_op::<IterateOnce>(definition.arguments),
             s => Ok(Box::new(UnimplementedOperation {
                 _name: s.to_string(),
             }) as Box<dyn TestOperation>),
@@ -2188,13 +2189,7 @@ impl TestOperation for IterateUntilDocumentOrError {
         async move {
             // A `SessionCursor` also requires a `&mut Session`, which would cause conflicting
             // borrows, so take the cursor from the map and return it after execution instead.
-            let mut cursor = test_runner
-                .entities
-                .write()
-                .await
-                .remove(id)
-                .unwrap()
-                .into_cursor();
+            let mut cursor = test_runner.take_cursor(id).await;
             let next = match &mut cursor {
                 TestCursor::Normal(cursor) => {
                     let mut cursor = cursor.lock().await;
@@ -2224,11 +2219,7 @@ impl TestOperation for IterateUntilDocumentOrError {
                 }
                 TestCursor::Closed => None,
             };
-            test_runner
-                .entities
-                .write()
-                .await
-                .insert(id.to_string(), Entity::Cursor(cursor));
+            test_runner.return_cursor(id, cursor).await;
             next.transpose()
                 .map(|opt| opt.map(|doc| Entity::Bson(Bson::Document(doc))))
         }
@@ -2965,6 +2956,49 @@ impl TestOperation for Upload {
     }
 }
 
+#[derive(Debug, Deserialize)]
+#[serde(rename_all = "camelCase", deny_unknown_fields)]
+pub(super) struct IterateOnce {}
+
+impl TestOperation for IterateOnce {
+    fn execute_entity_operation<'a>(
+        &'a self,
+        id: &'a str,
+        test_runner: &'a TestRunner,
+    ) -> BoxFuture<'a, Result<Option<Entity>>> {
+        async move {
+            let mut cursor = test_runner.take_cursor(id).await;
+            match &mut cursor {
+                TestCursor::Normal(cursor) => {
+                    let mut cursor = cursor.lock().await;
+                    cursor.try_advance().await?;
+                }
+                TestCursor::Session { cursor, session_id } => {
+                    cursor
+                        .try_advance(
+                            test_runner
+                                .entities
+                                .write()
+                                .await
+                                .get_mut(session_id)
+                                .unwrap()
+                                .as_mut_session_entity(),
+                        )
+                        .await?;
+                }
+                TestCursor::ChangeStream(change_stream) => {
+                    let mut change_stream = change_stream.lock().await;
+                    change_stream.next_if_any().await?;
+                }
+                TestCursor::Closed => panic!("Attempted to call IterateOnce on a closed cursor"),
+            }
+            test_runner.return_cursor(id, cursor).await;
+            Ok(None)
+        }
+        .boxed()
+    }
+}
+
 #[derive(Debug, Deserialize)]
 pub(super) struct UnimplementedOperation {
     _name: String,

src/test/spec/unified_runner/test_runner.rs

@@ -50,6 +50,7 @@ use super::{
     CollectionData,
     Entity,
     SessionEntity,
+    TestCursor,
     TestFileEntity,
 };
 
@@ -698,6 +699,27 @@ impl TestRunner {
             .as_client_encryption()
             .clone()
     }
+
+    /// Removes the cursor with the given ID from the entity map. This method passes ownership of
+    /// the cursor to the caller so that a mutable reference to a ClientSession can be accessed from
+    /// the entity map simultaneously. Once the caller is finished with the cursor, it MUST be
+    /// returned to the test runner via the return_cursor method below.
+    pub(crate) async fn take_cursor(&self, id: impl AsRef<str>) -> TestCursor {
+        self.entities
+            .write()
+            .await
+            .remove(id.as_ref())
+            .unwrap()
+            .into_cursor()
+    }
+
+    /// Returns the given cursor to the entity map. This method must be called after take_cursor.
+    pub(crate) async fn return_cursor(&self, id: impl AsRef<str>, cursor: TestCursor) {
+        self.entities
+            .write()
+            .await
+            .insert(id.as_ref().into(), Entity::Cursor(cursor));
+    }
 }
 
 #[cfg(feature = "in-use-encryption-unstable")]