{-
  Copyright (c) Meta Platforms, Inc. and affiliates.
  All rights reserved.

  This source code is licensed under the BSD-style license found in the
  LICENSE file in the root directory of this source tree.
-}

-- | Reexports 'Control.Concurrent.STM'
--   adding logging for deadlocks and utility functions
module Util.STM
 ( updateTVar
 , atomicallyWithLabel
 , Util.STM.atomically
 , module Control.Concurrent.STM
 ) where

import qualified Control.Concurrent.STM as STM
import Control.Concurrent.STM hiding (atomically)
import Control.Exception (Exception, BlockedIndefinitelyOnSTM, catch, throwIO)
import GHC.Stack
import Data.Typeable (Typeable)

-- | version of @modifyTVar'@ that returns the new value
updateTVar :: TVar a -> (a -> a) -> STM a
updateTVar var f = do
  x <- readTVar var
  let !x' = f x
  writeTVar var x'
  return x'

-- | Version of 'atomically' that takes a label
atomicallyWithLabel :: String -> STM a -> IO a
atomicallyWithLabel label stm = do
  STM.atomically stm `catch` (\(_::BlockedIndefinitelyOnSTM) ->
    throwIO (BlockedIndefinitelyOnNamedSTM label))

-- | Version of 'atomically' that labels the transaction using the call site
atomically :: HasCallStack => STM a -> IO a
atomically =
  atomicallyWithLabel (getCaller $ getCallStack callStack)

-- | The calling point is actually right next to the head as the head is the
-- current function (in this case it would be logInfo/Warning/Error/Fatal).
getCaller :: [(String, SrcLoc)] -> String
getCaller cs =
  case cs of
    -- Take the first element of the stack if it exists
    ((_,sl):_) -> prettySrcLoc sl
    _       -> "Unknown stack trace"

-- | 'atomicallyWithLabel' replaces occurrences of 'BlockedIndefinitelyOnSTM' w/
-- 'BlockedIndefinitelyOnNamedSTM', carrying the name of the transaction and
-- thus giving more helpful error messages.
newtype BlockedIndefinitelyOnNamedSTM = BlockedIndefinitelyOnNamedSTM String
    deriving (Typeable)

instance Show BlockedIndefinitelyOnNamedSTM where
    showsPrec _ (BlockedIndefinitelyOnNamedSTM name) =
        showString $ "thread blocked indefinitely in STM transaction: " ++ name

instance Exception BlockedIndefinitelyOnNamedSTM