RUST-1046 Fix iteration of cursors when batchSize doesn't divide result size (#484)

Author: patrickfreed

src/cursor/session.rs

@@ -52,7 +52,6 @@ pub struct SessionCursor<T>
 where
     T: DeserializeOwned + Unpin,
 {
-    exhausted: bool,
     client: Client,
     info: CursorInformation,
     buffer: VecDeque<T>,
@@ -63,10 +62,7 @@ where
     T: DeserializeOwned + Unpin + Send + Sync,
 {
     pub(crate) fn new(client: Client, spec: CursorSpecification<T>) -> Self {
-        let exhausted = spec.id() == 0;
-
         Self {
-            exhausted,
             client,
             info: spec.info,
             buffer: spec.initial_buffer,
@@ -165,12 +161,25 @@ where
     }
 }
 
+impl<T> SessionCursor<T>
+where
+    T: DeserializeOwned + Unpin,
+{
+    fn mark_exhausted(&mut self) {
+        self.info.id = 0;
+    }
+
+    fn is_exhausted(&self) -> bool {
+        self.info.id == 0
+    }
+}
+
 impl<T> Drop for SessionCursor<T>
 where
     T: DeserializeOwned + Unpin,
 {
     fn drop(&mut self) {
-        if self.exhausted {
+        if self.is_exhausted() {
             return;
         }
 
@@ -220,7 +229,9 @@ where
     fn drop(&mut self) {
         // Update the parent cursor's state based on any iteration performed on this handle.
         self.session_cursor.buffer = self.generic_cursor.take_buffer();
-        self.session_cursor.exhausted = self.generic_cursor.is_exhausted();
+        if self.generic_cursor.is_exhausted() {
+            self.session_cursor.mark_exhausted();
+        }
     }
 }
 

src/test/coll.rs

@@ -1049,3 +1049,46 @@ async fn collection_generic_bounds() {
         .collection(function_name!());
     let _result = coll.insert_one(Bar {}, None).await;
 }
+
+/// Verify that a cursor with multiple batches whose last batch isn't full
+/// iterates without errors.
+#[cfg_attr(feature = "tokio-runtime", tokio::test)]
+#[cfg_attr(feature = "async-std-runtime", async_std::test)]
+async fn cursor_batch_size() {
+    let _guard: RwLockReadGuard<()> = LOCK.run_concurrently().await;
+
+    let client = TestClient::new().await;
+    let coll = client
+        .init_db_and_coll("cursor_batch_size", "cursor_batch_size")
+        .await;
+
+    let doc = Document::new();
+    coll.insert_many(vec![&doc; 10], None).await.unwrap();
+
+    let opts = FindOptions::builder().batch_size(3).build();
+    let cursor_no_session = coll.find(doc! {}, opts.clone()).await.unwrap();
+    let docs: Vec<_> = cursor_no_session.try_collect().await.unwrap();
+    assert_eq!(docs.len(), 10);
+
+    // test session cursors
+    if client.is_standalone() {
+        return;
+    }
+    let mut session = client.start_session(None).await.unwrap();
+    let mut cursor = coll
+        .find_with_session(doc! {}, opts.clone(), &mut session)
+        .await
+        .unwrap();
+    let mut docs = Vec::new();
+    while let Some(doc) = cursor.next(&mut session).await {
+        docs.push(doc.unwrap());
+    }
+    assert_eq!(docs.len(), 10);
+
+    let mut cursor = coll
+        .find_with_session(doc! {}, opts, &mut session)
+        .await
+        .unwrap();
+    let docs: Vec<_> = cursor.stream(&mut session).try_collect().await.unwrap();
+    assert_eq!(docs.len(), 10);
+}