From fa95b65fada6602600f0059f3b7db8737aad1e30 Mon Sep 17 00:00:00 2001
From: "Michael S. Craig" <mkscrg@gmail.com>
Date: Mon, 12 Dec 2011 10:31:14 -0500
Subject: [PATCH] Fix for bad behavior when using TailableCursor.

`Database.MongoDB.Query.next` and `nextBatch` prefetch a promise of the next
batch of documents from the server whenever the current batch has been
exhausted. The following call to `next` or `nextBatch` fulfills that promise
(thereby turning it into a concrete batch of documents) and then returns one or
more documents in the batch.

The old behavior was to raise an exception if an empty batch with a nonzero
cursor ID was encountered. This is normal when using tailable cursors, so a
change was required.

Now, `Nothing` is returned with the still-live cursor ID, instead of raising
the exception. Also, prefetched empty batches with nonzero cursor IDs are
refetched once per call to avoid stale data.

This new prefetching behavior does not affect the performance of `next`, except
when calling it repeatedly on a tailable cursor with no new data. In those
(generally avoidable) cases, each call to `next` results in two server calls
instead of one.
---
 Database/MongoDB/Query.hs | 18 +++++++++++++++---
 1 file changed, 15 insertions(+), 3 deletions(-)

diff --git a/Database/MongoDB/Query.hs b/Database/MongoDB/Query.hs
index 3944692..bac7ff3 100644
--- a/Database/MongoDB/Query.hs
+++ b/Database/MongoDB/Query.hs
@@ -496,10 +496,16 @@ nextBatch :: (MonadIO m, MonadBaseControl IO m) => Cursor -> Action m [Document]
 -- ^ Return next batch of documents in query result, which will be empty if finished.
 nextBatch (Cursor fcol batchSize var) = modifyMVar var $ \dBatch -> do
 	-- Pre-fetch next batch promise from server and return current batch.
-	Batch limit cid docs <- fulfill dBatch
+	-- Discard pre-fetched batch if empty with nonzero cid.
+	Batch limit cid docs <- fulfill' dBatch
 	dBatch' <- if cid /= 0 then nextBatch' limit cid else return $ return (Batch 0 0 [])
 	return (dBatch', docs)
  where
+	fulfill' dBatch = do
+		b@(Batch limit cid docs) <- fulfill dBatch
+		if cid /= 0 && null docs
+			then nextBatch' limit cid >>= fulfill
+			else return b
 	nextBatch' limit cid = request [] (GetMore fcol batchSize' cid, remLimit)
 		where (batchSize', remLimit) = batchSizeRemainingLimit batchSize limit
 
@@ -507,9 +513,10 @@ next :: (MonadIO m, MonadBaseControl IO m) => Cursor -> Action m (Maybe Document
 -- ^ Return next document in query result, or Nothing if finished.
 next (Cursor fcol batchSize var) = modifyMVar var nextState where
 	-- Pre-fetch next batch promise from server when last one in current batch is returned.
+	-- Discard pre-fetched batch if empty with nonzero cid.
 	-- nextState:: DelayedBatch -> Action m (DelayedBatch, Maybe Document)
 	nextState dBatch = do
-		Batch limit cid docs <- fulfill dBatch
+		Batch limit cid docs <- fulfill' dBatch
 		case docs of
 			doc : docs' -> do
 				dBatch' <- if null docs' && cid /= 0
@@ -518,7 +525,12 @@ next (Cursor fcol batchSize var) = modifyMVar var nextState where
 				return (dBatch', Just doc)
 			[] -> if cid == 0
 				then return (return $ Batch 0 0 [], Nothing)  -- finished
-				else error $ "server returned empty batch but says more results on server"
+				else fmap (,Nothing) $ nextBatch' limit cid
+	fulfill' dBatch = do
+		b@(Batch limit cid docs) <- fulfill dBatch
+		if cid /= 0 && null docs
+			then nextBatch' limit cid >>= fulfill
+			else return b
 	nextBatch' limit cid = request [] (GetMore fcol batchSize' cid, remLimit)
 		where (batchSize', remLimit) = batchSizeRemainingLimit batchSize limit