warp: fix batch pump stalls

jsaddle-warp/jsaddle-warp.cabal

@@ -56,8 +56,10 @@ library
             Language.Javascript.JSaddle.WebSockets
         other-modules:
             Language.Javascript.JSaddle.WebSockets.Compat
+            Language.Javascript.JSaddle.WebSockets.BatchSender
         build-depends:
             aeson >=0.8.0.2 && <2.3,
+            async >=2.2 && <2.3,
             bytestring >=0.10.6.0 && <0.13,
             containers >=0.5.6.2 && <0.9,
             foreign-store >=0.2 && <0.3,
@@ -150,3 +152,21 @@ test-suite spec
     , websockets
   other-modules:
   default-language: Haskell2010
+
+test-suite batch-sender
+  if impl(ghcjs -any) || arch(javascript)
+    buildable: False
+
+  type: exitcode-stdio-1.0
+  main-is: BatchSenderSpec.hs
+  other-modules:
+      Language.Javascript.JSaddle.WebSockets.BatchSender
+  hs-source-dirs: tests src
+  ghc-options: -Wall -threaded
+  build-depends:
+      base
+    , async >=2.2 && <2.3
+    , hspec == 2.*
+    , jsaddle
+    , stm
+  default-language: Haskell2010

jsaddle-warp/src/Language/Javascript/JSaddle/WebSockets.hs

@@ -71,6 +71,7 @@ import Foreign.Store (readStore, lookupStore, writeStore, Store(..))
 import Language.Javascript.JSaddle (askJSM)
 import Control.Monad.IO.Class (MonadIO(..))
 
+import Language.Javascript.JSaddle.WebSockets.BatchSender (withBatchSender)
 import Language.Javascript.JSaddle.WebSockets.Compat (getTextMessageByteString)
 import qualified Data.Text.Encoding as T (decodeUtf8)
 
@@ -83,8 +84,8 @@ jsaddleOr opts entryPoint otherApp = do
             conn <- acceptRequest pending_conn
             initMsg <- receiveDataMessage conn
             case getTextMessageByteString initMsg of
-                Just "" -> do
-                    rec (processResult, processSyncResult, start) <- runJavaScript (sendTextData conn . encode) $ do
+                Just "" -> withBatchSender (sendTextData conn . encode) $ \publishBatch -> do
+                    rec (processResult, processSyncResult, start) <- runJavaScript publishBatch $ do
                             syncKey <- T.pack . show . contextId <$> askJSM
                             liftIO $ atomicModifyIORef' syncHandlers (\m -> (M.insert syncKey processSyncResult m, ()))
                             liftIO $ atomicModifyIORef' asyncHandlers (\m -> (M.insert syncKey processResult m, ()))

jsaddle-warp/src/Language/Javascript/JSaddle/WebSockets/BatchSender.hs

@@ -0,0 +1,28 @@
+module Language.Javascript.JSaddle.WebSockets.BatchSender
+    ( withBatchSender
+    ) where
+
+import Control.Concurrent.Async (link, withAsync)
+import Control.Concurrent.STM
+    ( atomically, newEmptyTMVarIO, putTMVar, takeTMVar, tryTakeTMVar )
+import Control.Monad (forever, void)
+import Language.Javascript.JSaddle.Types (Batch)
+
+-- | Send batches without blocking the pump.
+-- Keep at most one active batch and one pending batch.
+-- Sender failure interrupts connection.
+-- On exit, cancel the sender and wait for it to stop.
+--
+-- This should be called only from the serial JSaddle batch pump.
+-- Publish the next batch only after acknowledgement of the previous batch.
+-- Replace only WebSocket copies of batches already delivered through HTTP.
+withBatchSender :: (Batch -> IO ()) -> ((Batch -> IO ()) -> IO a) -> IO a
+withBatchSender send action = do
+    pending <- newEmptyTMVarIO
+    let publish batch = atomically $ do
+            void $ tryTakeTMVar pending
+            putTMVar pending batch
+        sender = forever $ atomically (takeTMVar pending) >>= send
+    withAsync sender $ \worker -> do
+        link worker
+        action publish

jsaddle-warp/tests/BatchSenderSpec.hs

@@ -0,0 +1,99 @@
+module Main (main) where
+
+import Control.Concurrent.Async
+    ( ExceptionInLinkedThread(..), cancel, withAsync )
+import Control.Concurrent.MVar
+    ( newEmptyMVar, putMVar, takeMVar, tryTakeMVar )
+import Control.Concurrent.STM
+    ( atomically, newTQueueIO, readTQueue, writeTQueue )
+import Control.Exception (Exception, finally, fromException, throwIO)
+import Control.Monad (forM_)
+import Language.Javascript.JSaddle.Types (Batch(..))
+import Language.Javascript.JSaddle.WebSockets.BatchSender (withBatchSender)
+import System.Timeout (timeout)
+import Test.Hspec
+
+main :: IO ()
+main = hspec $ around_ within $ describe "batch sender" $ do
+    it "advances acknowledged batches while the WebSocket blocks" $ do
+        started <- newEmptyMVar
+        release <- newEmptyMVar
+        delivered <- newTQueueIO
+        let send (Batch _ _ n) = do
+                if n == 1
+                    then putMVar started () >> takeMVar release
+                    else return ()
+                atomically $ writeTQueue delivered n
+        withBatchSender send $ \publish -> do
+            publish $ batch 1
+            takeMVar started
+            -- These batches have HTTP acknowledgements.
+            -- Keep only the latest pending WebSocket copy.
+            forM_ [2..1000] $ publish . batch
+            putMVar release ()
+            first <- atomically $ readTQueue delivered
+            latest <- atomically $ readTQueue delivered
+            (first, latest) `shouldBe` (1, 1000)
+
+    it "interrupts the connection on sender failure" $ do
+        blocked <- newEmptyMVar
+        let send _ = throwIO SendFailure
+        (withBatchSender send $ \publish -> do
+            publish $ batch 1
+            takeMVar blocked) `shouldThrow` linkedSendFailure
+
+    it "stops the sender on normal connection exit" $ do
+        started <- newEmptyMVar
+        stopped <- newEmptyMVar
+        blocked <- newEmptyMVar
+        let send _ = (putMVar started () >> takeMVar blocked)
+                `finally` putMVar stopped ()
+        withBatchSender send $ \publish -> do
+            publish $ batch 1
+            takeMVar started
+        tryTakeMVar stopped `shouldReturn` Just ()
+
+    it "preserves connection errors and stops the sender" $ do
+        started <- newEmptyMVar
+        stopped <- newEmptyMVar
+        blocked <- newEmptyMVar
+        let send _ = (putMVar started () >> takeMVar blocked)
+                `finally` putMVar stopped ()
+        (withBatchSender send $ \publish -> do
+            publish $ batch 1
+            takeMVar started
+            throwIO ConnectionFailure) `shouldThrow` (== ConnectionFailure)
+        tryTakeMVar stopped `shouldReturn` Just ()
+
+    it "stops the sender when the connection is cancelled" $ do
+        started <- newEmptyMVar
+        stopped <- newEmptyMVar
+        blockedSend <- newEmptyMVar
+        blockedConnection <- newEmptyMVar
+        let send _ = (putMVar started () >> takeMVar blockedSend)
+                `finally` putMVar stopped ()
+        withAsync (withBatchSender send $ \publish -> do
+            publish $ batch 1
+            takeMVar blockedConnection) $ \connection -> do
+                takeMVar started
+                cancel connection
+                tryTakeMVar stopped `shouldReturn` Just ()
+
+batch :: Int -> Batch
+batch = Batch [] False
+
+data Failure = SendFailure | ConnectionFailure
+    deriving (Eq, Show)
+
+instance Exception Failure
+
+linkedSendFailure :: ExceptionInLinkedThread -> Bool
+linkedSendFailure (ExceptionInLinkedThread _ exception) =
+    fromException exception == Just SendFailure
+
+within :: IO () -> IO ()
+within action = do
+    result <- timeout 5000000 action
+    case result of
+        Nothing -> expectationFailure "batch sender test timed out"
+        Just () -> return ()