From 316ae02ecc0e76b62dd112c391013b0b62c75822 Mon Sep 17 00:00:00 2001
From: Victor Denisov <denisovenator@gmail.com>
Date: Sun, 23 Oct 2016 21:52:10 -0700
Subject: [PATCH] Handle errors properly for ordered updates

---
 Database/MongoDB/Query.hs | 91 ++++++++++++++++++++++-----------------
 1 file changed, 51 insertions(+), 40 deletions(-)

diff --git a/Database/MongoDB/Query.hs b/Database/MongoDB/Query.hs
index 140385f..612ff6d 100644
--- a/Database/MongoDB/Query.hs
+++ b/Database/MongoDB/Query.hs
@@ -645,17 +645,20 @@ update' ordered col updateDocs = do
                       updates
   let lens = map length chunks
   let lSums = 0 : (zipWith (+) lSums lens)
-  exceptionThrown <- liftIO $ newIORef False
-  blocks <- forM (zip lSums chunks) $ \b -> do
-    ctx <- ask
-    liftIO $ do
-      et <- readIORef exceptionThrown
-      if et && ordered
-        then return $ UpdateResult True 0 Nothing [] [] [] -- TODO probably should be revised
-        else runReaderT (updateBlock ordered col b) ctx
-      `catch` \(e :: SomeException) -> do
-        writeIORef exceptionThrown True
-        return $ UpdateResult True 0 Nothing [] [] [] -- TODO  probably should be revised
+  errorDetected <- liftIO $ newIORef False
+  ctx <- ask
+  blocks <- forM (zip lSums chunks) $ \b -> liftIO $ do
+    ed <- readIORef errorDetected
+    if ed && ordered
+      then return $ UpdateResult True 0 Nothing [] [] []
+      else do
+        ur <- runReaderT (updateBlock ordered col b) ctx
+        when (failed ur) $ do
+          writeIORef errorDetected True
+        return ur
+    `catch` \(e :: SomeException) -> do
+      writeIORef errorDetected True
+      return $ UpdateResult True 0 Nothing [] [] [] -- TODO  probably should be revised
   let failedTotal = or $ map failed blocks
   let updatedTotal = sum $ map nMatched blocks
   let modifiedTotal =
@@ -700,38 +703,46 @@ updateBlockLegacy :: (MonadIO m)
 updateBlockLegacy ordered col (prevCount, docs) = do
   db <- thisDatabase
   ctx <- ask
+  errorDetected <- liftIO $ newIORef False
   results <-
     liftIO $ forM (zip [prevCount, (prevCount + 1) ..] docs) $ \(i, updateDoc) -> do
-      let doc = (at "u" updateDoc) :: Document
-      let sel = (at "q" updateDoc) :: Document
-      let upsrt = if at "upsert" updateDoc then [Upsert] else []
-      let multi = if at "multi" updateDoc then [MultiUpdate] else []
-      mRes <- runReaderT (write (Update (db <.> col) (upsrt ++ multi) sel doc)) ctx
-      case mRes of
-        Nothing -> return $ UpdateResult False 0 Nothing [] [] []
-        Just resDoc -> do
-          let em = lookup "err" resDoc
-          let eCode = lookup "code" resDoc
-          let wtimeout = fromMaybe False $ lookup "wtimeout" resDoc
-          if isNothing em && isNothing eCode
-            then do
-              let n = at "n" resDoc
-              let ups = do
-                    upsValue <- lookup "upserted" resDoc
-                    return $ Upserted i upsValue
-              return $ UpdateResult False n Nothing (maybeToList ups) [] []
-            else do
-              let defaultCode = if wtimeout then 64 else 24
-              let errV = fromMaybe "unknown error" em
-              let c = fromMaybe defaultCode eCode
-              if wtimeout
+      ed <- readIORef errorDetected
+      if ed && ordered
+        then do
+          return $ UpdateResult True 0 Nothing [] [] []
+        else do
+          let doc = (at "u" updateDoc) :: Document
+          let sel = (at "q" updateDoc) :: Document
+          let upsrt = if at "upsert" updateDoc then [Upsert] else []
+          let multi = if at "multi" updateDoc then [MultiUpdate] else []
+          mRes <- runReaderT (write (Update (db <.> col) (upsrt ++ multi) sel doc)) ctx
+          case mRes of
+            Nothing -> return $ UpdateResult False 0 Nothing [] [] []
+            Just resDoc -> do
+              let em = lookup "err" resDoc
+              let eCode = lookup "code" resDoc
+              let wtimeout = fromMaybe False $ lookup "wtimeout" resDoc
+              if isNothing em && isNothing eCode
                 then do
-                  return $ UpdateResult True 0 Nothing [] [] [WriteConcernError c errV]
+                  let n = at "n" resDoc
+                  let ups = do
+                        upsValue <- lookup "upserted" resDoc
+                        return $ Upserted i upsValue
+                  return $ UpdateResult False n Nothing (maybeToList ups) [] []
                 else do
-                  return $ UpdateResult True 0 Nothing [] [WriteError i c errV] []
-      `catch` \(e :: SomeException) -> do
-                 when ordered $ liftIO $ throwIO e
-                 return $ UpdateResult True 0 Nothing [] [WriteError i 0 (show e)] []
+                  let defaultCode = if wtimeout then 64 else 24
+                  let errV = fromMaybe "unknown error" em
+                  let c = fromMaybe defaultCode eCode
+                  if wtimeout
+                    then do
+                      writeIORef errorDetected True
+                      return $ UpdateResult True 0 Nothing [] [] [WriteConcernError c errV]
+                    else do
+                      writeIORef errorDetected True
+                      return $ UpdateResult True 0 Nothing [] [WriteError i c errV] []
+          `catch` \(e :: SomeException) -> do
+                     writeIORef errorDetected True
+                     return $ UpdateResult True 0 Nothing [] [WriteError i 0 (show e)] []
   return $ foldl1' mergeUpdateResults results
 
 mergeUpdateResults :: UpdateResult -> UpdateResult -> UpdateResult
@@ -859,7 +870,7 @@ deleteBlock ordered col docs = do
         liftIO $ forM docs $ \deleteDoc -> do
           let sel = (at "q" deleteDoc) :: Document
           let opts = if at "limit" deleteDoc == (1 :: Int) then [SingleRemove] else []
-          runReaderT (write (Delete (db <.> col) opts sel)) ctx
+          _ <- runReaderT (write (Delete (db <.> col) opts sel)) ctx
           return Nothing
           `catch` \(e :: SomeException) -> do
                      when ordered $ liftIO $ throwIO e