module GHC.Stack.Profiler.Internal.Manager (
  Manager (..),
  newManager,
  stopManager,
  shouldProfile,
  enableEventLogging,
  disableEventLogging,
  enableSampling,
  disableSampling,
  registerSamplerThread,
  unregisterSamplerThread,
  stopAllSamplerThreads,

  -- * Sampler Threads
  Sampler (..),
  cancelSampler,

  -- * Event Loop
  EventLoop (..),
  startEventLoop,
  stopEventLoop,

  -- * Control Messages
  ControlMessage (..),
  startProfiling,
  stopProfiling,
  sendPublishInitEventMessages,
  sendStartProfilingMessage,
  sendStopProfilingMessage,
  sendEnableEventlogMessage,
  sendDisableEventlogMessage,
) where

import Control.Concurrent (ThreadId)
import Control.Concurrent.Async (Async (..), async, cancel, link)
import Control.Concurrent.Chan
import Control.Concurrent.MVar
import Control.Concurrent.STM (STM)
import Control.Concurrent.STM.TVar
import qualified Control.Concurrent.STM.TVar as STM
import qualified Control.Concurrent.STM.TVar as TVar
import Control.Monad (forever)
import Control.Monad.STM (atomically)
import Data.ByteString (ByteString)
import qualified Data.ByteString.Lazy as BSL
import Data.Foldable (for_)
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import qualified Debug.Trace
import qualified Debug.Trace.Binary.Compat as Compat
import GHC.Generics (Generic)
import qualified GHC.Stack.Profiler.Core as GSPC (Message (ProtocolVersion), ProtocolVersion (MyProtocolVersion))
import qualified GHC.Stack.Profiler.Internal.Decode as Decode
import GHC.Stack.Profiler.Internal.SymbolTable

-- NOTE: The `Manager` type (but not its implementation) is part of the public API.

-- | A `Manager` handle, which can be used to stop the manager with `stopManager`.
--
--   @since 0.5.0.0
data Manager = MkManager
  { Manager -> TVar (Map ThreadId Sampler)
samplerThreadMapVar :: !(TVar (Map ThreadId Sampler))
  -- ^ 'Async' of the stack sampling thread.
  , Manager -> TVar (Maybe EventLoop)
eventLoopThreadVar :: !(TVar (Maybe EventLoop))
  -- ^ Main event loop thread responsible for processing profiler messages, etc...
  , Manager -> StackSymbolTable
symbolTableRef :: !StackSymbolTable
  -- ^ Global table for common symbols.
  , Manager -> TVar Bool
shouldSampleVar :: !(TVar Bool)
  -- ^ Is the profiler currently running?
  --
  -- Can be controlled via 'startProfiler' and 'stopProfiler'.
  -- This variable describes whether the user wants to profile, regardless
  -- of the eventlog state.
  , Manager -> TVar Bool
eventLoggingStartedVar :: !(TVar Bool)
  -- ^ Is there an eventlog?
  --
  -- It is fully possible that we start profiling but no eventlog-writer
  -- being connected/configured. The eventlog can be enabled at a later point,
  -- or stopped/started via @eventlog-socket@.
  -- This variable tracks the state of the eventlog-writer.
  , Manager -> Chan ControlMessage
messageChan :: Chan ControlMessage
  }
  deriving ((forall x. Manager -> Rep Manager x)
-> (forall x. Rep Manager x -> Manager) -> Generic Manager
forall x. Rep Manager x -> Manager
forall x. Manager -> Rep Manager x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
$cfrom :: forall x. Manager -> Rep Manager x
from :: forall x. Manager -> Rep Manager x
$cto :: forall x. Rep Manager x -> Manager
to :: forall x. Rep Manager x -> Manager
Generic, Manager -> Manager -> Bool
(Manager -> Manager -> Bool)
-> (Manager -> Manager -> Bool) -> Eq Manager
forall a. (a -> a -> Bool) -> (a -> a -> Bool) -> Eq a
$c== :: Manager -> Manager -> Bool
== :: Manager -> Manager -> Bool
$c/= :: Manager -> Manager -> Bool
/= :: Manager -> Manager -> Bool
Eq)

newManager :: Bool -> IO Manager
newManager :: Bool -> IO Manager
newManager Bool
wait = do
  tracingEnabled <- IO Bool
Compat.userTracingEnabledIO
  samplerThreadMapVar <- newTVarIO Map.empty
  eventLoopThreadVar <- newTVarIO Nothing
  symbolTableRef <- emptySymbolTableIO
  shouldSampleVar <- newTVarIO (not wait)
  eventLoggingStartedVar <- newTVarIO tracingEnabled
  messageChan <- newChan
  pure
    MkManager
      { samplerThreadMapVar
      , eventLoopThreadVar
      , symbolTableRef
      , shouldSampleVar
      , eventLoggingStartedVar
      , messageChan
      }

-- NOTE: `stopManager` is part of the public API.

-- | Stop a `Manager`.
--
--   This also stops every `Sampler` started by this manager.
--
--   __Warning:__ If the `Manager` is not stopped before the program exits,
--   some messages may not be written to the eventlog.
--
--   @since 0.5.0.0
stopManager :: Manager -> IO ()
stopManager :: Manager -> IO ()
stopManager Manager
manager = do
  Manager -> IO ()
stopAllSamplerThreads Manager
manager
  Manager -> IO ()
stopEventLoop Manager
manager

-- | Can we profile right now?
--
-- We only sample a stack if the profiler is instructed to run and the eventlog is enabled.
shouldProfile :: Manager -> STM Bool
shouldProfile :: Manager -> STM Bool
shouldProfile Manager
manager =
  (Bool -> Bool -> Bool) -> STM Bool -> STM Bool -> STM Bool
forall a b c. (a -> b -> c) -> STM a -> STM b -> STM c
forall (f :: * -> *) a b c.
Applicative f =>
(a -> b -> c) -> f a -> f b -> f c
liftA2
    Bool -> Bool -> Bool
(&&)
    (TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar (TVar Bool -> STM Bool) -> TVar Bool -> STM Bool
forall a b. (a -> b) -> a -> b
$ Manager -> TVar Bool
shouldSampleVar Manager
manager)
    (TVar Bool -> STM Bool
forall a. TVar a -> STM a
readTVar (TVar Bool -> STM Bool) -> TVar Bool -> STM Bool
forall a b. (a -> b) -> a -> b
$ Manager -> TVar Bool
eventLoggingStartedVar Manager
manager)

enableEventLogging :: Manager -> STM ()
enableEventLogging :: Manager -> STM ()
enableEventLogging Manager
manager = do
  TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
TVar.writeTVar (Manager -> TVar Bool
eventLoggingStartedVar Manager
manager) Bool
True

disableEventLogging :: Manager -> STM ()
disableEventLogging :: Manager -> STM ()
disableEventLogging Manager
manager = do
  TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
TVar.writeTVar (Manager -> TVar Bool
eventLoggingStartedVar Manager
manager) Bool
False

enableSampling :: Manager -> STM ()
enableSampling :: Manager -> STM ()
enableSampling Manager
manager = do
  TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
TVar.writeTVar (Manager -> TVar Bool
shouldSampleVar Manager
manager) Bool
True

disableSampling :: Manager -> STM ()
disableSampling :: Manager -> STM ()
disableSampling Manager
manager = do
  TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
TVar.writeTVar (Manager -> TVar Bool
shouldSampleVar Manager
manager) Bool
False

registerSamplerThread :: Manager -> Sampler -> IO ()
registerSamplerThread :: Manager -> Sampler -> IO ()
registerSamplerThread Manager
manager samplerThread :: Sampler
samplerThread@MkSampler{Async ()
samplerAsync :: Async ()
samplerAsync :: Sampler -> Async ()
samplerAsync} = do
  Async () -> IO ()
forall a. Async a -> IO ()
link Async ()
samplerAsync -- If the sampler crashes, we want to know.
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    TVar (Map ThreadId Sampler)
-> (Map ThreadId Sampler -> Map ThreadId Sampler) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
STM.modifyTVar' (Manager -> TVar (Map ThreadId Sampler)
samplerThreadMapVar Manager
manager) ((Map ThreadId Sampler -> Map ThreadId Sampler) -> STM ())
-> (Map ThreadId Sampler -> Map ThreadId Sampler) -> STM ()
forall a b. (a -> b) -> a -> b
$ \Map ThreadId Sampler
threadMap ->
      (ThreadId -> Sampler -> Map ThreadId Sampler -> Map ThreadId Sampler
forall k a. Ord k => k -> a -> Map k a -> Map k a
Map.insert (Async () -> ThreadId
forall a. Async a -> ThreadId
asyncThreadId Async ()
samplerAsync) Sampler
samplerThread Map ThreadId Sampler
threadMap)

unregisterSamplerThread :: Manager -> Sampler -> IO ()
unregisterSamplerThread :: Manager -> Sampler -> IO ()
unregisterSamplerThread Manager
manager MkSampler{Async ()
samplerAsync :: Sampler -> Async ()
samplerAsync :: Async ()
samplerAsync} =
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ do
    TVar (Map ThreadId Sampler)
-> (Map ThreadId Sampler -> Map ThreadId Sampler) -> STM ()
forall a. TVar a -> (a -> a) -> STM ()
STM.modifyTVar'
      (Manager -> TVar (Map ThreadId Sampler)
samplerThreadMapVar Manager
manager)
      (ThreadId -> Map ThreadId Sampler -> Map ThreadId Sampler
forall k a. Ord k => k -> Map k a -> Map k a
Map.delete (Async () -> ThreadId
forall a. Async a -> ThreadId
asyncThreadId Async ()
samplerAsync))

stopAllSamplerThreads :: Manager -> IO ()
stopAllSamplerThreads :: Manager -> IO ()
stopAllSamplerThreads Manager
manager = do
  samplerThreads <-
    STM [Sampler] -> IO [Sampler]
forall a. STM a -> IO a
atomically (STM [Sampler] -> IO [Sampler]) -> STM [Sampler] -> IO [Sampler]
forall a b. (a -> b) -> a -> b
$ do
      samplerThreadMap <- TVar (Map ThreadId Sampler) -> STM (Map ThreadId Sampler)
forall a. TVar a -> STM a
readTVar (Manager -> TVar (Map ThreadId Sampler)
samplerThreadMapVar Manager
manager)
      writeTVar (samplerThreadMapVar manager) Map.empty
      pure $ Map.elems samplerThreadMap
  for_ samplerThreads cancelSampler

-------------------------------------------------------------------------------
-- Sampler Threads
-------------------------------------------------------------------------------

-- NOTE: The `Sampler` type (but not its implementation) is part of the public API.

-- | A `Sampler` handle, which can be used to stop the sampler with `GHC.Stack.Profiler.stopSampler`.
--
--   @since 0.5.0.0
newtype Sampler = MkSampler
  { Sampler -> Async ()
samplerAsync :: Async ()
  }

cancelSampler :: Sampler -> IO ()
cancelSampler :: Sampler -> IO ()
cancelSampler MkSampler{Async ()
samplerAsync :: Sampler -> Async ()
samplerAsync :: Async ()
samplerAsync} =
  Async () -> IO ()
forall a. Async a -> IO ()
cancel Async ()
samplerAsync

-------------------------------------------------------------------------------
-- Event Loop
-------------------------------------------------------------------------------

newtype EventLoop = MkEventLoop
  { EventLoop -> Async ()
eventLoopAsync :: Async ()
  }

data ControlMessage
  = WriteProfileSample [ByteString]
  | PublishInitEvents (MVar ())
  | StartProfiling (MVar ())
  | StopProfiling (MVar ())
  | StartEventlog (MVar ())
  | StopEventlog (MVar ())

startEventLoop :: Manager -> IO ()
startEventLoop :: Manager -> IO ()
startEventLoop Manager
manager = do
  !eventLoopThread <- do
    eventLoopAsync <- IO () -> IO (Async ())
forall a. IO a -> IO (Async a)
async (IO () -> IO (Async ())) -> IO () -> IO (Async ())
forall a b. (a -> b) -> a -> b
$ IO () -> IO ()
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO ()) -> IO () -> IO ()
forall a b. (a -> b) -> a -> b
$ Manager -> IO ()
eventHandler Manager
manager
    link eventLoopAsync -- If the event loop crashes, we want to know.
    pure $ MkEventLoop{eventLoopAsync}
  atomically $ do
    writeTVar (eventLoopThreadVar manager) (Just eventLoopThread)

eventHandler :: Manager -> IO ()
eventHandler :: Manager -> IO ()
eventHandler Manager
manager = do
  msg <- Chan ControlMessage -> IO ControlMessage
forall a. Chan a -> IO a
readChan (Manager -> Chan ControlMessage
messageChan Manager
manager)
  run <- atomically $ shouldProfile manager
  case msg of
    WriteProfileSample [ByteString]
msgs ->
      case Bool
run of
        Bool
True ->
          (ByteString -> IO ()) -> [ByteString] -> IO ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
(a -> m b) -> t a -> m ()
mapM_ ByteString -> IO ()
Compat.traceBinaryEventIO [ByteString]
msgs
        Bool
False ->
          -- If we received a sample but the eventlog is currently locked
          -- discard the message.
          () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    StartProfiling MVar ()
barrier -> do
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ Manager -> STM ()
enableSampling Manager
manager
      MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
barrier ()
    StopProfiling MVar ()
barrier -> do
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ Manager -> STM ()
disableSampling Manager
manager
      MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
barrier ()
    StartEventlog MVar ()
barrier -> do
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ Manager -> STM ()
enableEventLogging Manager
manager
      MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
barrier ()
    StopEventlog MVar ()
barrier -> do
      STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ Manager -> STM ()
disableEventLogging Manager
manager
      MVar () -> () -> IO ()
forall a. MVar a -> a -> IO ()
putMVar MVar ()
barrier ()
    PublishInitEvents MVar ()
barrier -> do
      symbolTable <- STM (SymbolTableWriter MapTable) -> IO (SymbolTableWriter MapTable)
forall a. STM a -> IO a
atomically (STM (SymbolTableWriter MapTable)
 -> IO (SymbolTableWriter MapTable))
-> STM (SymbolTableWriter MapTable)
-> IO (SymbolTableWriter MapTable)
forall a b. (a -> b) -> a -> b
$ StackSymbolTable -> STM (SymbolTableWriter MapTable)
readSymbolTable (Manager -> StackSymbolTable
symbolTableRef Manager
manager)
      let
        versionMessage = ProtocolVersion -> Message
GSPC.ProtocolVersion ProtocolVersion
GSPC.MyProtocolVersion
        messages = Message
versionMessage Message -> [Message] -> [Message]
forall a. a -> [a] -> [a]
: SymbolTableWriter MapTable -> [Message]
Decode.initMessages SymbolTableWriter MapTable
symbolTable
        messagesBytes = [Message] -> [ByteString]
Decode.serializeMessages [Message]
messages

      for_ messagesBytes $ \ByteString
binaryMessage ->
        ByteString -> IO ()
Compat.traceBinaryEventIO (ByteString -> ByteString
BSL.toStrict ByteString
binaryMessage)

      Debug.Trace.flushEventLog
      putMVar barrier ()

stopEventLoop :: Manager -> IO ()
stopEventLoop :: Manager -> IO ()
stopEventLoop Manager
manager = do
  maybeEventThread <- STM (Maybe EventLoop) -> IO (Maybe EventLoop)
forall a. STM a -> IO a
atomically (STM (Maybe EventLoop) -> IO (Maybe EventLoop))
-> STM (Maybe EventLoop) -> IO (Maybe EventLoop)
forall a b. (a -> b) -> a -> b
$ TVar (Maybe EventLoop)
-> (Maybe EventLoop -> (Maybe EventLoop, Maybe EventLoop))
-> STM (Maybe EventLoop)
forall s a. TVar s -> (s -> (a, s)) -> STM a
stateTVar (Manager -> TVar (Maybe EventLoop)
eventLoopThreadVar Manager
manager) (,Maybe EventLoop
forall a. Maybe a
Nothing)
  case maybeEventThread of
    Maybe EventLoop
Nothing ->
      -- Manager is already stopped
      () -> IO ()
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure ()
    Just MkEventLoop{Async ()
eventLoopAsync :: EventLoop -> Async ()
eventLoopAsync :: Async ()
eventLoopAsync} -> do
      -- Send stopProfilingMessage. This tells the event loop to flush all messages.
      Manager -> IO ()
sendStopProfilingMessage Manager
manager
      -- Stop the event loop thread.
      Async () -> IO ()
forall a. Async a -> IO ()
cancel Async ()
eventLoopAsync

-------------------------------------------------------------------------------
-- Events
-------------------------------------------------------------------------------

-- NOTE: The `startProfiling` function is part of the public API.

-- | Start all `Sampler` threads.
--
--   This blocks until all `Sampler` threads have started.
--
--   __Warning:__ This function deadlocks when used with a stopped `Manager`.
--
--  @since 0.5.0.0
startProfiling :: Manager -> IO ()
startProfiling :: Manager -> IO ()
startProfiling Manager
manager = do
  -- TODO: This atomically is redundant, the main loop thread sets it anyway.
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (Manager -> TVar Bool
shouldSampleVar Manager
manager) Bool
True
  Manager -> IO ()
sendStartProfilingMessage Manager
manager

-- NOTE: The `stopProfiling` function is part of the public API.

-- | Start all `Sampler` threads.
--
--   This blocks until all `Sampler` threads have stopped.
--
--   __Warning:__ This function deadlocks when used with a stopped `Manager`.
--
--  @since 0.5.0.0
stopProfiling :: Manager -> IO ()
stopProfiling :: Manager -> IO ()
stopProfiling Manager
manager = do
  -- TODO: This atomically is *not* redundant. It makes sure no new samples
  -- can be created. Otherwise, new samples could be created and queued while
  -- we are waiting for the event loop to process this message. It is
  -- important that, once this message is processed, no sampler thread is
  -- sampling at all. Otherwise, there will be new init events that are not
  -- published.
  STM () -> IO ()
forall a. STM a -> IO a
atomically (STM () -> IO ()) -> STM () -> IO ()
forall a b. (a -> b) -> a -> b
$ TVar Bool -> Bool -> STM ()
forall a. TVar a -> a -> STM ()
writeTVar (Manager -> TVar Bool
shouldSampleVar Manager
manager) Bool
False
  Manager -> IO ()
sendStopProfilingMessage Manager
manager

-- | Start profiling.
--
-- Blocks until the message has been processed by the main event loop.
sendStartProfilingMessage :: Manager -> IO ()
sendStartProfilingMessage :: Manager -> IO ()
sendStartProfilingMessage Manager
manager = do
  barrier <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  writeChan
    (messageChan manager)
    (StartProfiling barrier)
  takeMVar barrier

-- | Stop profiling.
--
-- Blocks until the message has been processed by the main event loop.
sendStopProfilingMessage :: Manager -> IO ()
sendStopProfilingMessage :: Manager -> IO ()
sendStopProfilingMessage Manager
manager = do
  barrier <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  writeChan
    (messageChan manager)
    (StopProfiling barrier)
  takeMVar barrier

-- | Start EventLogging now.
--
-- Blocks until the message has been processed by the main event loop.
sendEnableEventlogMessage :: Manager -> IO ()
sendEnableEventlogMessage :: Manager -> IO ()
sendEnableEventlogMessage Manager
manager = do
  barrier <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  writeChan
    (messageChan manager)
    (StartEventlog barrier)
  takeMVar barrier

-- | Stop EventLogging now.
--
-- Blocks until the message has been processed by the main event loop.
sendDisableEventlogMessage :: Manager -> IO ()
sendDisableEventlogMessage :: Manager -> IO ()
sendDisableEventlogMessage Manager
manager = do
  barrier <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  writeChan
    (messageChan manager)
    (StopEventlog barrier)
  takeMVar barrier

-- | Publish all init messages so far.
--
-- Blocks until the init events have been written to the eventlog and
-- eventlog was flushed.
sendPublishInitEventMessages :: Manager -> IO ()
sendPublishInitEventMessages :: Manager -> IO ()
sendPublishInitEventMessages Manager
manager = do
  barrier <- IO (MVar ())
forall a. IO (MVar a)
newEmptyMVar
  writeChan
    (messageChan manager)
    (PublishInitEvents barrier)
  takeMVar barrier