-- SPDX-FileCopyrightText: 2024 Steffen Jost <s.jost@fraport.de>
--
-- SPDX-License-Identifier: AGPL-3.0-or-later

module Handler.Utils.Concurrent
  ( module Handler.Utils.Concurrent  
  ) where

-- NOTE: use `retrySTM` and `checkSTM` instead of `retry` or `check`

import Import
import UnliftIO.Concurrent as Handler.Utils.Concurrent hiding (yield)



-- | Run a handler action until it finishes or if it exceeds a given number of microseconds via `registerDelay`
timeoutHandler :: Int -> HandlerFor site a -> HandlerFor site (Maybe a)
timeoutHandler maxWait act = do
    innerAct <- handlerToIO
    (hresult, tid) <- liftIO $ do
      hresult <- newTVarIO Nothing
      tid <- forkIO $ do
        res <- innerAct act
        atomically $ writeTVar hresult $ Just res
      return (hresult, tid)
    res <- liftIO $ do
      flag <- registerDelay maxWait
      atomically $ do
        out <- readTVar flag
        res <- readTVar hresult
        checkSTM $ out || isJust res
        return res
    case res of
      Nothing -> liftIO $ do
        killThread tid
        readTVarIO hresult     -- read once more after kill to ensure that any result is noticed
      _ -> return res