STM debouncer

ghcide/src/Development/IDE/Core/Debouncer.hs

@@ -9,12 +9,11 @@ module Development.IDE.Core.Debouncer
     ) where
 
 import           Control.Concurrent.Async
-import           Control.Concurrent.STM.Stats (atomically, atomicallyNamed)
-import           Control.Exception
-import           Control.Monad                (join)
-import           Data.Foldable                (traverse_)
+import           Control.Concurrent.STM
+import           Control.Concurrent.STM.Stats (atomicallyNamed)
+import           Control.Monad                (join, void)
 import           Data.Hashable
-import qualified Focus
+import           GHC.Conc                     (unsafeIOToSTM)
 import qualified StmContainers.Map            as STM
 import           System.Time.Extra
 
@@ -35,24 +34,50 @@ newAsyncDebouncer = Debouncer . asyncRegisterEvent <$> STM.newIO
 
 -- | Register an event that will fire after the given delay if no other event
 -- for the same key gets registered until then.
---
--- If there is a pending event for the same key, the pending event will be killed.
--- Events are run unmasked so it is up to the user of `registerEvent`
--- to mask if required.
-asyncRegisterEvent :: (Eq k, Hashable k) => STM.Map k (Async ()) -> Seconds -> k -> IO () -> IO ()
-asyncRegisterEvent d 0 k fire = do
-    join $ atomically $ do
-        prev <- STM.focus Focus.lookupAndDelete k d
-        return $ traverse_ cancel prev
-    fire
-asyncRegisterEvent d delay k fire = mask_ $ do
-    a <- asyncWithUnmask $ \unmask -> unmask $ do
-        sleep delay
-        fire
-        atomically $ STM.delete k d
-    do
-        prev <- atomicallyNamed "debouncer" $ STM.focus (Focus.lookup <* Focus.insert a) k d
-        traverse_ cancel prev
+asyncRegisterEvent
+    :: (Eq k, Hashable k)
+    => STM.Map k (TVar (Maybe (Seconds, IO())))
+    -> Seconds
+    -> k
+    -> IO ()
+    -> IO ()
+asyncRegisterEvent d delay k fire = join $ atomicallyNamed "debouncer - register" $ do
+    -- The previous TVar for this key, if any
+    prev <- STM.lookup k d
+    case prev of
+        Just v -> do
+            current <- readTVar v
+            case current of
+                -- Not empty, means that there is a thread running the actions
+                Just _  -> writeTVar v (Just (delay, fire)) >> return (pure ())
+                -- Empty = no thread. We need to start one for running the action
+                Nothing -> writeTVar v (Just (delay, fire)) >> return (restart v)
+
+        -- No previous TVar, we need to insert one and restart a thread for running the action
+        Nothing
+          | delay == 0 -> return fire
+          | otherwise  -> do
+            var <- newTVar (Just (delay, fire))
+            STM.insert var k d
+            return (restart var)
+    where
+        -- | Restart a thread to run the action stored in the given TVar
+        --   Once the action is done, the thread dies.
+        --   Assumes the Tvar is not empty
+        restart var =
+            void $ async $
+                join $ atomicallyNamed "debouncer - sleep" $ do
+                    contents <- readTVar var
+                    case contents of
+                        Nothing -> error "impossible"
+                        Just (s,act) -> do
+                            -- sleep for the given delay
+                            -- If the TVar is written while sleeping,
+                            -- the transaction will restart
+                            unsafeIOToSTM $ sleep s
+                            -- we are done - empty the TVar before exiting
+                            writeTVar var Nothing
+                            return act
 
 -- | Debouncer used in the DAML CLI compiler that emits events immediately.
 noopDebouncer :: Debouncer k