Fix an unbounded blocked-thread leak in runJavaScript; ObjC over-retains + jsaddleWebViewInvalidate for wkwebview

jsaddle-wkwebview/cbits-cocoa/AppDelegate.m

@@ -57,6 +57,12 @@ -(void)applicationWillFinishLaunching:(NSNotification *)notification {
         jsaddleCallbacks.callWithWebView((void *)webView, _handler);
     if (jsaddleCallbacks.callIO)
         jsaddleCallbacks.callIO(global_willFinishLaunchingWithOptions);
+    // Balance the allocs above: -setContentView: retained the webview, so the
+    // window is its sole owner from here.  Without these releases the boot
+    // webview — and its WebContent renderer process — outlives the window
+    // forever.  After the callbacks, so nothing reads webView post-release.
+    [webView release];
+    [theConfiguration release];
     // Listen to URL events
     NSAppleEventManager *manager = [NSAppleEventManager sharedAppleEventManager];
     [manager

jsaddle-wkwebview/cbits/WKWebView-cbits.m

@@ -54,26 +54,33 @@ - (instancetype)initHandler:(HsStablePtr)startHandler resultHandler:(HsStablePtr
     return self;
 }
 
+// The stable-pointer checks in the three handler methods below cover the
+// window between removeJSaddleHandler freeing them (setting them to NULL) and
+// the last WebKit delivery already in flight: dereferencing a freed
+// StablePtr is undefined, ignoring a message for a detached context is not.
 - (void)webView:(WKWebView *)webView didFinishNavigation:(WKNavigation *)navigation
 {
-    if (jsaddleCallbacks.jsaddleStart)
+    if (jsaddleCallbacks.jsaddleStart && self.startHandler)
         jsaddleCallbacks.jsaddleStart(self.startHandler);
 }
 
 - (void)userContentController:(WKUserContentController *)userContentController didReceiveScriptMessage:(WKScriptMessage *)message
 {
     NSString * s = (NSString *) message.body;
-    if (jsaddleCallbacks.jsaddleResult)
+    if (jsaddleCallbacks.jsaddleResult && self.resultHandler)
         jsaddleCallbacks.jsaddleResult(self.resultHandler, [s UTF8String]);
 }
 
 - (void)webView:(WKWebView *)webView runJavaScriptTextInputPanelWithPrompt:(NSString *)prompt
     defaultText:(NSString *)defaultText initiatedByFrame:(WKFrameInfo *)frame completionHandler:(void (^)(NSString *result))completionHandler
 {
     if([prompt isEqualToString:@"JSaddleSync"]) {
-        _completionHandler = completionHandler;
-        if (jsaddleCallbacks.jsaddleSyncResult)
+        if (jsaddleCallbacks.jsaddleSyncResult && self.syncHandler) {
+            _completionHandler = completionHandler;
             jsaddleCallbacks.jsaddleSyncResult(self.syncHandler, (void *)self, [defaultText UTF8String]);
+        }
+        else
+            completionHandler(@"");   // detached context: unblock the page
     }
     else {
         // TODO decide what if anything should go here
@@ -96,14 +103,55 @@ - (void)webView:(WKWebView *)webView decidePolicyForNavigationAction:(WKNavigati
 
 @end
 
-void addJSaddleHandler(WKWebView *webView, HsStablePtr startHandler, HsStablePtr resultHandler, HsStablePtr syncHandler) {
+// Returns the handler (retained: the caller's reference, released by
+// removeJSaddleHandler) so a host that tears a webview down can detach it
+// again — see removeJSaddleHandler.
+JSaddleHandler * addJSaddleHandler(WKWebView *webView, HsStablePtr startHandler, HsStablePtr resultHandler, HsStablePtr syncHandler) {
     JSaddleHandler * handler = [[JSaddleHandler alloc] initHandler:startHandler resultHandler:resultHandler syncHandler:syncHandler];
     [[[webView configuration] userContentController] addScriptMessageHandler:handler name:@"jsaddle"];
 #ifdef USE_UIKIT
     webView.scrollView.bounces = NO;
 #endif
     webView.navigationDelegate = handler;
     webView.UIDelegate = handler;
+    return handler;
+}
+
+// Detach a handler installed by addJSaddleHandler and free the three stable
+// pointers it holds.  Those pointers keep their closures — and through them the
+// context's whole callback graph (for a reflex application, the entire live
+// network) — alive for the life of the PROCESS, which in ghci mode means every
+// reloaded instance's network stays reachable: measured at ~30MB of live heap
+// leaked per reload.
+//
+// On the MAIN QUEUE, because that is where all three callbacks are delivered
+// from: running the detach there is what makes it mutually exclusive with a
+// callback that is mid-flight, so the frees cannot race a deRefStablePtr.
+// Asynchronously, because the caller may be a Haskell thread that the main
+// thread is (or is about to be) waiting on — a dispatch_sync from the jsaddle
+// sync round trip's own producer would deadlock.  The block retains both
+// objects, so the host may release the webview as soon as this returns.
+// @onDetached@ is run (once) after the handler is detached; freeing the three
+// stable pointers is its job, not this file's.  Calling the RTS directly here
+// is not an option: in ghci mode this file is compiled into a plain dylib with
+// no RTS to link against, so hs_free_stable_ptr is an undefined symbol — the
+// same constraint that makes every Haskell call out of this file go through the
+// registered callback table (see WKWebView-callbacks.h).
+void removeJSaddleHandler(WKWebView *webView, JSaddleHandler *handler, HsStablePtr onDetached) {
+    dispatch_async(dispatch_get_main_queue(), ^{
+        [[[webView configuration] userContentController] removeScriptMessageHandlerForName:@"jsaddle"];
+        if (webView.navigationDelegate == handler) webView.navigationDelegate = nil;
+        if (webView.UIDelegate == handler) webView.UIDelegate = nil;
+        // Clear before the pointers are freed, so a delivery WebKit had already
+        // queued behind this block sees a detached handler (the methods above
+        // check) rather than a dangling stable pointer.
+        handler.startHandler = NULL;
+        handler.resultHandler = NULL;
+        handler.syncHandler = NULL;
+        [handler release];   // balances addJSaddleHandler's +1
+        if (jsaddleCallbacks.callIO)
+            jsaddleCallbacks.callIO(onDetached);
+    });
 }
 
 void evaluateJavaScript(WKWebView *webView, const char * _Nonnull js) {
@@ -118,11 +166,12 @@ void completeSync(JSaddleHandler *handler, const char * _Nonnull s) {
     NSString *string = [[NSString alloc] initWithCString:s encoding:NSUTF8StringEncoding];
     handler.completionHandler(string);
     handler.completionHandler = NULL;
+    [string release];
 }
 
 void loadHTMLStringWithBaseURL(WKWebView *webView, const char * _Nonnull html, const char * _Nonnull url) {
   NSString *htmlString = [NSString stringWithCString:html encoding:NSUTF8StringEncoding];
-  NSURL *baseURL = [NSURL URLWithString:[[NSString alloc] initWithCString:url encoding:NSUTF8StringEncoding]];
+  NSURL *baseURL = [NSURL URLWithString:[NSString stringWithCString:url encoding:NSUTF8StringEncoding]];
   dispatch_async(dispatch_get_main_queue(), ^{
       [webView loadHTMLString:htmlString baseURL:baseURL];
     });

jsaddle-wkwebview/src-ghc/Language/Javascript/JSaddle/WKWebView.hs

@@ -3,6 +3,7 @@ module Language.Javascript.JSaddle.WKWebView
     , jsaddleMainHTMLWithBaseURL
     , jsaddleMainFile
     , WKWebView(..)
+    , jsaddleWebViewInvalidate
     , run
     , run'
     , runWithAppConfig
@@ -23,7 +24,8 @@ import Data.Word
 import Foreign.Marshal.Utils (fromBool)
 import Language.Javascript.JSaddle.WKWebView.Internal
        (jsaddleMain, jsaddleMainHTMLWithBaseURL, jsaddleMainFile, WKWebView(..),
-        mainBundleResourcePath, registerJSaddleCallbacks)
+        mainBundleResourcePath, registerJSaddleCallbacks,
+        jsaddleWebViewInvalidate)
 import System.Environment (getProgName)
 import Foreign.C.String (CString, withCString)
 import Foreign.C.Types (CBool, CInt)

jsaddle-wkwebview/src/Language/Javascript/JSaddle/WKWebView/Internal.hs

@@ -1,17 +1,22 @@
-{-# LANGUAGE ForeignFunctionInterface, OverloadedStrings #-}
+{-# LANGUAGE ForeignFunctionInterface, OverloadedStrings, ScopedTypeVariables #-}
 module Language.Javascript.JSaddle.WKWebView.Internal
     ( jsaddleMain
     , jsaddleMainHTMLWithBaseURL
     , jsaddleMainFile
     , WKWebView(..)
     , mainBundleResourcePath
     , registerJSaddleCallbacks
+    , jsaddleWebViewInvalidate
     ) where
 
 import Control.Monad (void, join, unless)
 import Control.Concurrent (forkIO, forkOS)
+import Control.Exception (IOException, handle)
 import Control.Concurrent.MVar
-       (MVar, modifyMVar_, newEmptyMVar, newMVar, putMVar, takeMVar)
+       (MVar, modifyMVar, modifyMVar_, newEmptyMVar, newMVar, putMVar, readMVar,
+        takeMVar, tryPutMVar, withMVar)
+
+import Data.IORef (newIORef, readIORef, writeIORef, atomicModifyIORef')
 
 import System.IO.Unsafe (unsafePerformIO)
 
@@ -26,9 +31,11 @@ import Data.Text.Encoding (encodeUtf8)
 import Foreign.C.String (CString)
 import Foreign.C.Types (CBool(..), CInt(..))
 import Foreign.Ptr (Ptr, FunPtr, nullPtr)
-import Foreign.StablePtr (StablePtr, newStablePtr, deRefStablePtr)
+import Foreign.StablePtr
+       (StablePtr, newStablePtr, deRefStablePtr, freeStablePtr)
 
-import Language.Javascript.JSaddle (Results, Batch, JSM)
+import Language.Javascript.JSaddle (Results, JSM)
+import Language.Javascript.JSaddle.Types (Batch(..))
 import Language.Javascript.JSaddle.Run (runJavaScriptWithSerializer)
 import Language.Javascript.JSaddle.Run.Files (initState, runBatch, ghcjsHelpers)
 
@@ -47,6 +54,42 @@ newtype JSaddleHandler = JSaddleHandler (Ptr JSaddleHandler)
 wkWebViewBatchLock :: MVar ()
 wkWebViewBatchLock = unsafePerformIO (newMVar ())
 
+-- | Per-context invalidators, keyed by the raw webview pointer each context
+-- sends its batches to.  'jsaddleWebViewInvalidate' runs (and drops) a
+-- webview's entries so the host application can safely release the webview:
+-- until then jsaddle's threads keep calling 'evaluateJavaScript' with the
+-- stored pointer, and releasing the webview under them is a use-after-free.
+{-# NOINLINE contextInvalidators #-}
+contextInvalidators :: MVar [(Ptr WKWebView, IO ())]
+contextInvalidators = unsafePerformIO (newMVar [])
+
+-- | Invalidate every jsaddle context attached to @webView@: no jsaddle thread
+-- will touch the webview pointer after this returns, so the host may release
+-- the webview (and let its WebContent renderer process exit).  Blocks until
+-- any in-flight batch send has completed.  The contexts' Haskell threads are
+-- not stopped — their subsequent sends throw instead (a dropped send would
+-- leave jsaddle's batch thread waiting for a reply forever while holding
+-- 'wkWebViewBatchLock', deadlocking every other window; see 'jsaddleMain'').
+--
+-- Also detaches each context's ObjC handler and frees the stable pointers it
+-- holds, which is what makes an invalidated context COLLECTABLE — see
+-- @removeJSaddleHandler@ in cbits/WKWebView-cbits.m.  Without it the callback
+-- closures, and through them the context's entire object graph, are pinned for
+-- the life of the process.
+--
+-- Call it once per webview: the registry entries are claimed under the lock, so
+-- a second concurrent call for the SAME webview finds nothing and returns while
+-- the first is still running (it would not wait for that in-flight send).
+jsaddleWebViewInvalidate :: WKWebView -> IO ()
+jsaddleWebViewInvalidate (WKWebView ptr) = do
+    -- CLAIM under the lock, RUN outside it.  An invalidator blocks until this
+    -- context's in-flight send finishes; holding the registry across that would
+    -- stall every other webview's invalidation — and every new context's
+    -- registration, which 'jsaddleMain'' does on the same MVar during startup.
+    invs <- modifyMVar contextInvalidators $ \invs ->
+        return (filter ((/= ptr) . fst) invs, [ inv | (p, inv) <- invs, p == ptr ])
+    sequence_ invs
+
 foreign export ccall jsaddleStart :: StablePtr (IO ()) -> IO ()
 foreign export ccall jsaddleResult :: StablePtr (Results -> IO ()) -> CString -> IO ()
 foreign export ccall jsaddleSyncResult :: StablePtr (Results -> IO Batch) -> JSaddleHandler -> CString -> IO ()
@@ -116,7 +159,8 @@ registerJSaddleCallbacks = modifyMVar_ callbacksRegistered $ \done -> do
         c_jsaddleWkSetCallbacks start result sync io cstr cstrb webv cint
     return True
 
-foreign import ccall addJSaddleHandler :: WKWebView -> StablePtr (IO ()) -> StablePtr (Results -> IO ()) -> StablePtr (Results -> IO Batch) -> IO ()
+foreign import ccall addJSaddleHandler :: WKWebView -> StablePtr (IO ()) -> StablePtr (Results -> IO ()) -> StablePtr (Results -> IO Batch) -> IO JSaddleHandler
+foreign import ccall removeJSaddleHandler :: WKWebView -> JSaddleHandler -> StablePtr (IO ()) -> IO ()
 foreign import ccall loadHTMLStringWithBaseURL :: WKWebView -> CString -> CString -> IO ()
 foreign import ccall loadBundleFile :: WKWebView -> CString -> CString -> IO ()
 foreign import ccall evaluateJavaScript :: WKWebView -> CString -> IO ()
@@ -155,23 +199,107 @@ jsaddleMainFile url allowing f webView =
                 loadBundleFile webView u a
 
 jsaddleMain' :: JSM () -> WKWebView -> IO () -> IO ()
-jsaddleMain' f webView loadHtml = do
+jsaddleMain' f webView@(WKWebView ptr) loadHtml = do
     registerJSaddleCallbacks
     ready <- newEmptyMVar
 
+    -- Every use of the raw webview pointer goes through this gate.  An MVar
+    -- (not an IORef) so 'jsaddleWebViewInvalidate' SYNCHRONIZES with an
+    -- in-flight send: invalidation blocks until the send's evaluateJavaScript
+    -- call — which copies the pointer into a main-queue block, retaining it —
+    -- has returned, and no later send can start.  After that the host can
+    -- release the webview with no thread left to dereference it.
+    aliveVar <- newMVar True
+    -- This context's pending detach (see @removeJSaddleHandler@), from the
+    -- moment 'addJSaddleHandler' has installed something to detach until
+    -- invalidation runs it.  Taken and written under @aliveVar@ — by the
+    -- invalidator below and by the install site — so the detach happens exactly
+    -- once even if the context is invalidated while it is still being set up.
+    detachRef <- newIORef Nothing
+    modifyMVar_ contextInvalidators $ return .
+        ((ptr, do
+             modifyMVar_ aliveVar $ \_ -> do
+                 mbDetach <- atomicModifyIORef' detachRef (\d -> (Nothing, d))
+                 sequence_ mbDetach
+                 return False
+             -- Release the startup thread below if the page never signalled
+             -- ready (invalidated mid-load): nothing else will ever fill this,
+             -- and it is holding an OS thread and this context's closures.  It
+             -- wakes, hits the now-closed gate and exits.
+             void $ tryPutMVar ready ()) :)
+    -- An invalidated context must FAIL its send, not silently drop it.
+    -- 'runJavaScriptWithSerializer' holds the shared 'wkWebViewBatchLock'
+    -- across @sendBatch@ AND the @takeResult@ that waits for that batch's
+    -- reply: a dropped batch leaves its batch thread waiting for a reply that
+    -- can never arrive, holding the lock every OTHER window needs — so one
+    -- invalidated context deadlocks the whole process (in ghci mode: every
+    -- later @:main@ hangs on its first round trip, CPU idle).  Throwing
+    -- unwinds out of 'withMVar', releasing both MVars, and kills only the dead
+    -- context's own batch pump.
+    let withWebView act = withMVar aliveVar $ \alive ->
+            if alive then act
+            else ioError (userError "jsaddle-wkwebview: context invalidated \
+                                    \(webview released) — batch dropped")
+
     (processResult, syncResult, start) <- runJavaScriptWithSerializer (Just wkWebViewBatchLock) (\batch ->
-        useAsCString (toStrict $ "runJSaddleBatch(" <> encode batch <> ");") $
-            evaluateJavaScript webView)
+        withWebView $
+            useAsCString (toStrict $ "runJSaddleBatch(" <> encode batch <> ");") $
+                evaluateJavaScript webView)
         f
 
+    -- INBOUND must be gated too, or invalidating a context leaks threads without
+    -- bound.  The webview keeps its script-message handler until the host
+    -- releases it, and a page whose JS is still running keeps posting results;
+    -- 'runJavaScriptWithSerializer' forks a @processResults@ thread per delivery,
+    -- and with this context's batch pump gone (killed by the send gate above)
+    -- every one of them blocks forever on the pump's MVar.  Measured before this
+    -- gate: ~5000 blocked threads PER SECOND, 500k threads and 12GB RSS in one
+    -- session (and 72GB in a longer one).  A dead context simply drops what
+    -- arrives: nothing is listening for it anyway.
+    --
+    -- The flag is only READ here (not held across the handler): holding it would
+    -- deadlock the sync path, whose 'readMVar lastAsyncBatch' can wait on the
+    -- batch thread — which needs this same MVar to send.
+    let ifAlive :: a -> IO a -> IO a
+        ifAlive dflt act = readMVar aliveVar >>= \alive ->
+            if alive then act else return dflt
     startHandler <- newStablePtr (putMVar ready ())
-    resultHandler <- newStablePtr processResult
-    syncResultHandler <- newStablePtr syncResult
-    addJSaddleHandler webView startHandler resultHandler syncResultHandler
+    resultHandler <- newStablePtr (\results -> ifAlive () (processResult results))
+    -- An empty batch is what a context with nothing to say returns.
+    syncResultHandler <- newStablePtr
+        (\results -> ifAlive (Batch [] False 0) (syncResult results))
+    handler <- addJSaddleHandler webView startHandler resultHandler syncResultHandler
+    -- Freeing the three stable pointers above is what lets an invalidated
+    -- context be collected: each pins its closure, and those closures capture
+    -- the whole 'JSContextRef' — for a reflex application, the entire live
+    -- network.  Held for the life of the process (as they were), every instance
+    -- a ghci session ever ran stays reachable: ~30MB of live heap per reload.
+    -- Only safe once the ObjC handler can no longer dereference them, so it is
+    -- @removeJSaddleHandler@ that runs this, after detaching (asynchronously, on
+    -- the main queue).  @selfRef@ lets the callback free its own stable pointer
+    -- too; it is written before any detach can be requested below.
+    selfRef <- newIORef Nothing
+    onDetached <- newStablePtr $ do
+        freeStablePtr startHandler
+        freeStablePtr resultHandler
+        freeStablePtr syncResultHandler
+        readIORef selfRef >>= mapM_ freeStablePtr
+    writeIORef selfRef (Just onDetached)
+    -- Arm the detach for the invalidator — or, if invalidation already happened
+    -- (it found nothing to detach), detach right here.
+    let detach = removeJSaddleHandler webView handler onDetached
+    modifyMVar_ aliveVar $ \alive -> do
+        if alive then writeIORef detachRef (Just detach) else detach
+        return alive
     loadHtml
-    void . forkOS $ do
+    -- If the context is invalidated before the page signals ready, the gate
+    -- throws here; that just means this window died during startup, so swallow
+    -- it rather than letting forkOS report an unhandled exception.  Only that
+    -- 'IOException' — anything else is a real failure and belongs on stderr.
+    void . forkOS . handle (\(_ :: IOException) -> return ()) $ do
         takeMVar ready
-        useAsCString (toStrict jsaddleJs) (evaluateJavaScript webView) >> void (forkIO start)
+        withWebView $ useAsCString (toStrict jsaddleJs) (evaluateJavaScript webView)
+        void (forkIO start)
 
 jsaddleStart :: StablePtr (IO ()) -> IO ()
 jsaddleStart ptrHandler = join (deRefStablePtr ptrHandler)