Repository: mongodb/mongo-rust-driver (language: rust)

Issue:
`aggregate_with_session` fails when collection size > batchSize
## Overview

Instantiate a MongoDB DB with a collection of N documents. Run a simple aggregation against that collection with a `batchSize` that is lower than N and that's not a multiple of N.

Examples:
- `N` = 100, `batchSize` = 40
- `N` = 102, `batchSize` = 101 (default value)


For all those cases, `cursor.next(&mut session)` will return the following error:

```
Error: Command failed (CursorNotFound): cursor id XXXXXXXXX not found)
```

## Reproduction

Below is a minimal reproduction example:

```rs
use anyhow::Result;
use mongodb::{
    bson::{doc, Document},
    options::{AggregateOptions, ClientOptions},
    Client, Collection,
};
use tokio;

#[tokio::main]
async fn main() -> Result<()> {
    let client_options = ClientOptions::parse("MONGO_URL").await?;
    let client = Client::with_options(client_options)?;
    let mut session = client.start_session(None).await?;
    let db = client.database("DATABASE_NAME");
    let coll: Collection<Document> = db.collection("COLLECTION_NAME");
    // Default batchSize is 101. Have at least 102 documents for the error to happen
    // or set a `batchSize` according to the rules above depending on the number of documents in your collection
    let opts = AggregateOptions::builder().build();
    let mut res = coll
        .aggregate_with_session(
            vec![doc! {
                "$project": { "_id": 1 }
            }],
            opts,
            &mut session,
        )
        .await?;
    let mut docs = vec![];

    while let Some(result) = res.next(&mut session).await {
        match result {
            Ok(document) => docs.push(document),
            Err(e) => return Err(e.into()),
        }
    }

    println!("Number of documents: {}", docs.len());

    Ok(())
}
```


Discussion / hints:

Fix patch:
diff --git a/src/cursor/session.rs b/src/cursor/session.rs
index e8634513d..a64977857 100644
--- a/src/cursor/session.rs
+++ b/src/cursor/session.rs
@@ -61,7 +61,6 @@ pub struct SessionCursor<T>
 where
     T: DeserializeOwned + Unpin,
 {
-    exhausted: bool,
     client: Client,
     info: CursorInformation,
     buffer: VecDeque<T>,
@@ -79,10 +78,7 @@ where
         spec: CursorSpecification<T>,
         pinned: Option<PinnedConnectionHandle>,
     ) -> Self {
-        let exhausted = spec.id() == 0;
-
         Self {
-            exhausted,
             client,
             info: spec.info,
             buffer: spec.initial_buffer,
@@ -198,12 +194,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;
         }
 
@@ -254,7 +263,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();
+        }
     }
 }
 


Test patch:
diff --git a/src/test/coll.rs b/src/test/coll.rs
index 4f9f1b4c1..4de2b9070 100644
--- a/src/test/coll.rs
+++ b/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);
+}