module GHC.Stack.Profiler (
  -- * High-Level API

  -- ** Profiler
  Profiler (..),
  withProfiler,
  withProfilerWith,
  withProfilerFromEnv,
  startProfiler,
  startProfilerWith,
  startProfilerFromEnv,
  stopProfiler,

  -- ** Options
  Options (
    wait,
    shouldSample,
    sampleRtsThreads,
    sampleProfilerThreads,
    sampleInterval
  ),
  defaultOptions,
  Interval (..),

  -- *** Thread Filters and Glob Patterns
  ThreadFilter,
  ThreadLabel,
  ShouldSample (..),
  Glob,
  matches,
  sampleInclude,
  sampleExclude,
  sampleIncludeExclude,

  -- *** Environment Variables
  fromEnv,

  -- * Low-Level API

  -- ** Manager
  Manager,
  withManager,
  startManager,
  stopManager,

  -- ** Commands
  startProfiling,
  stopProfiling,

  -- ** Samplers
  Sampler,
  withSamplerForMe,
  startSamplerFor,
  startSamplerWith,
  stopSampler,
) where

import Control.Concurrent.Async (Async (..))
import Control.Exception
import Control.Monad.IO.Class (MonadIO (..))
import Data.Bifunctor (Bifunctor (..))
import Data.Foldable (traverse_)
import Data.Functor ((<&>))
import Data.IORef (IORef, newIORef, readIORef, writeIORef)
import Data.List (isPrefixOf)
import qualified Data.Map.Strict as Map
import Data.Maybe (catMaybes, fromMaybe)
import Data.Set (Set)
import qualified Data.Set as S
import qualified Data.Set as Set
import Data.String (IsString (..))
import GHC.Conc
import GHC.Conc.Sync (threadLabel)
import GHC.IsList (IsList (..))
import qualified GHC.Stack.Profiler.Internal.Eventlog.Socket as Eventlog.Socket
import GHC.Stack.Profiler.Internal.Manager
import GHC.Stack.Profiler.Internal.Sampler (Interval (MkIntervalMillis), SamplerDescr (MkSamplerDescr), startSampler, stopSampler, withSampler)
import qualified GHC.Stack.Profiler.Internal.Sampler as SamplerDescr
import GHC.Stack.Profiler.Internal.Util (DList, Glob, WriterT, matches, runWriterT, tell)
import System.Environment (lookupEnv)
import System.IO (hPutStrLn, stderr)
import Text.Printf (printf)
import Text.Read (readMaybe)

-------------------------------------------------------------------------------
-- High-level API
-------------------------------------------------------------------------------

-------------------------------------------------------------------------------
-- Profiler

-- | A profiler handle, which can be used to stop the profiler with `stopProfiler`.
--
--   @since 0.5.0.0
data Profiler = MkProfiler
  { Profiler -> Manager
profilerManager :: !Manager
  , Profiler -> Sampler
profilerSampler :: !Sampler
  }

-- | Run an action with a `Profiler` and the default `Options`.
--
--   __Warning:__ This function spawns a `Manager` thread.
--   Having multiple concurrent `Manager` threads is unsupported and unsafe.
--
--   @since 0.5.0.0
withProfiler :: (Profiler -> IO a) -> IO a
withProfiler :: forall a. (Profiler -> IO a) -> IO a
withProfiler Profiler -> IO a
action =
  IO Profiler -> (Profiler -> IO ()) -> (Profiler -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO Profiler
startProfiler Profiler -> IO ()
stopProfiler Profiler -> IO a
action

-- | Variant of `withProfiler` that accepts `Options`.
--
--   @since 0.5.0.0
withProfilerWith :: Options -> (Profiler -> IO a) -> IO a
withProfilerWith :: forall a. Options -> (Profiler -> IO a) -> IO a
withProfilerWith Options
options Profiler -> IO a
action =
  IO Profiler -> (Profiler -> IO ()) -> (Profiler -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Options -> IO Profiler
startProfilerWith Options
options) Profiler -> IO ()
stopProfiler Profiler -> IO a
action

-- | Variant of `withProfiler` that reads `Options` from the environment.
--
--   If @GHC_STACK_PROFILER@ is unset or empty, no `Profiler` is started.
--
--   @since 0.5.0.0
withProfilerFromEnv :: (Maybe Profiler -> IO a) -> IO a
withProfilerFromEnv :: forall a. (Maybe Profiler -> IO a) -> IO a
withProfilerFromEnv Maybe Profiler -> IO a
action =
  IO (Maybe Profiler)
-> (Maybe Profiler -> IO ()) -> (Maybe Profiler -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket IO (Maybe Profiler)
startProfilerFromEnv ((Profiler -> IO ()) -> Maybe Profiler -> IO ()
forall (t :: * -> *) (f :: * -> *) a b.
(Foldable t, Applicative f) =>
(a -> f b) -> t a -> f ()
traverse_ Profiler -> IO ()
stopProfiler) Maybe Profiler -> IO a
action

-- | Start a `Profiler` with the default `Options`.
--
--   This function returns a `Profiler` handle, which can be used to stop
--   the profiler with `stopProfiler`.
--
--   __Warning:__ This function spawns a `Manager` thread.
--   Having multiple concurrent `Manager` threads is unsupported and unsafe.
--
--   __Warning:__ If the `Profiler` is not stopped before the program exits,
--   some messages may not be written to the eventlog.
--
--   @since 0.5.0.0
startProfiler :: IO Profiler
startProfiler :: IO Profiler
startProfiler =
  Options -> IO Profiler
startProfilerWith Options
defaultOptions

-- | Variant of `startProfiler` that accepts `Options`.
--
--   @since 0.5.0.0
startProfilerWith :: Options -> IO Profiler
startProfilerWith :: Options -> IO Profiler
startProfilerWith Options
options = do
  profilerManager <- Bool -> IO Manager
startManager (Options -> Bool
wait Options
options)
  profilerSampler <- startSamplerWith profilerManager options
  pure MkProfiler{profilerManager, profilerSampler}

-- | Variant of `startProfiler` that accepts `Options`.
--
--   If @GHC_STACK_PROFILER@ is unset or empty, no `Profiler` is started.
--
--   @since 0.5.0.0
startProfilerFromEnv :: IO (Maybe Profiler)
startProfilerFromEnv :: IO (Maybe Profiler)
startProfilerFromEnv =
  IO (Maybe Options)
fromEnv IO (Maybe Options)
-> (Maybe Options -> IO (Maybe Profiler)) -> IO (Maybe Profiler)
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= (Options -> IO Profiler) -> Maybe Options -> IO (Maybe Profiler)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> Maybe a -> f (Maybe b)
traverse Options -> IO Profiler
startProfilerWith

-- | Stop a `Profiler`.
--
--   @since 0.5.0.0
stopProfiler :: Profiler -> IO ()
stopProfiler :: Profiler -> IO ()
stopProfiler MkProfiler{Manager
profilerManager :: Profiler -> Manager
profilerManager :: Manager
profilerManager, Sampler
profilerSampler :: Profiler -> Sampler
profilerSampler :: Sampler
profilerSampler} = do
  Manager -> Sampler -> IO ()
stopSampler Manager
profilerManager Sampler
profilerSampler
  Manager -> IO ()
stopManager Manager
profilerManager

-------------------------------------------------------------------------------
-- Options

-- | The options for `withProfilerWith` and `startProfilerWith`.
--
--   To construct options, modify `defaultOptions` using the fields:
--
--   [@`GHC.Stack.Profiler.wait` :: `Bool`@]:
--     Determines if sampler threads are started on creation or wait for a
--     "start profiling" command on the eventlog socket. If you are using
--     @ghc-stack-profiler@ with @eventlog-socket@'s control commands, this
--     should be set to @True@. Otherwise, this should be @False@. The default
--     is @False@.
--   [@`GHC.Stack.Profiler.shouldSample` :: `ThreadId` -> `Maybe` `ThreadLabel` -> `ShouldSample`@]:
--     Determines if the thread idenfied by the `ThreadId` should be sampled.
--     The current `ThreadLabel`, returned by `threadLabel`, is passed as the
--     second argument. If this function returns `Never`, the thread will never
--     be sampled, even if its `ThreadLabel` changes. The default predicate
--     always returns `Yes`. This function is not used for RTS threads or
--     threads spawned by @ghc-stack-profiler@.
--   [@`GHC.Stack.Profiler.sampleRtsThreads` :: `Bool`@]:
--     Determines if builtin RTS threads should be sampled. The builtin RTS
--     threads are the TimerManager and IOManager threads, and do not usually
--     have an interesting call-stack profile. The default is @False@.
--   [@`GHC.Stack.Profiler.sampleProfilerThreads` :: `Bool`@]:
--     Determines if the threads spawned by @ghc-stack-profiler@ should be
--     sampled. The default is @False@.
--   [@`GHC.Stack.Profiler.sampleInterval` :: `Interval`@]:
--     Determines the sampling interval.
--     The default is @10@ milliseconds.
--
--   @since 0.5.0.0
data Options = MkOptions
  { Options -> Bool
wait :: !Bool
  , Options -> ThreadFilter
shouldSample :: ThreadFilter
  , Options -> Bool
sampleRtsThreads :: !Bool
  , Options -> Bool
sampleProfilerThreads :: !Bool
  , Options -> Interval
sampleInterval :: !Interval
  }

-- | The default `Options`. See `Options` for the default values.
--
--   @since 0.5.0.0
defaultOptions :: Options
defaultOptions :: Options
defaultOptions =
  MkOptions
    { wait :: Bool
wait = Bool
False
    , shouldSample :: ThreadFilter
shouldSample = \ThreadId
_threadId Maybe ThreadLabel
_maybeThreadLabel -> ShouldSample
Yes
    , sampleRtsThreads :: Bool
sampleRtsThreads = Bool
False
    , sampleProfilerThreads :: Bool
sampleProfilerThreads = Bool
False
    , sampleInterval :: Interval
sampleInterval = Int -> Interval
MkIntervalMillis Int
10
    }

-- | A thread filter, used to determine which threads should be sampled.
--
--   Used in the `shouldSample` field of `Options`.
--
--   @since 0.5.0.0
type ThreadFilter = ThreadId -> Maybe ThreadLabel -> ShouldSample

-- | A thread label, as set by `labelThread`.
--
--   @since 0.5.0.0
type ThreadLabel = String

-- | The result type of a `ThreadFilter`.
--
--   @since 0.5.0.0
data ShouldSample
  = -- | The thread should be sampled.
    Yes
  | -- | The thread should not be sampled.
    No
  | -- | The thread should never be sampled.
    Never

-- | Construct a thread filter from an include `Glob` pattern.
--
--   If the thread label matches the given pattern, the thread filter returns `Yes`.
--   Otherwise, the thread filter returns `No`.
--   The thread filter never returns `Never`.
--
--   @since 0.5.0.0
sampleInclude ::
  -- | The include pattern.
  Glob ->
  ThreadFilter
sampleInclude :: Glob -> ThreadFilter
sampleInclude Glob
globInclude =
  (Maybe ThreadLabel -> ShouldSample) -> ThreadFilter
forall a b. a -> b -> a
const ((Maybe ThreadLabel -> ShouldSample) -> ThreadFilter)
-> ((ThreadLabel -> ShouldSample)
    -> Maybe ThreadLabel -> ShouldSample)
-> (ThreadLabel -> ShouldSample)
-> ThreadFilter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShouldSample
-> (ThreadLabel -> ShouldSample)
-> Maybe ThreadLabel
-> ShouldSample
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ShouldSample
No ((ThreadLabel -> ShouldSample) -> ThreadFilter)
-> (ThreadLabel -> ShouldSample) -> ThreadFilter
forall a b. (a -> b) -> a -> b
$
    Bool -> ShouldSample
fromBool (Bool -> ShouldSample)
-> (ThreadLabel -> Bool) -> ThreadLabel -> ShouldSample
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \ThreadLabel
label ->
      Glob
globInclude Glob -> ThreadLabel -> Bool
`matches` ThreadLabel
label

-- | Construct a thread filter from an exclude `Glob` pattern.
--
--   If the thread label matches the given pattern, the thread filter returns `No`.
--   Otherwise, the thread filter returns `Yes`.
--   The thread filter never returns `Never`.
--
--   @since 0.5.0.0
sampleExclude ::
  -- | The exclude pattern.
  Glob ->
  ThreadFilter
sampleExclude :: Glob -> ThreadFilter
sampleExclude Glob
globExclude =
  (Maybe ThreadLabel -> ShouldSample) -> ThreadFilter
forall a b. a -> b -> a
const ((Maybe ThreadLabel -> ShouldSample) -> ThreadFilter)
-> ((ThreadLabel -> ShouldSample)
    -> Maybe ThreadLabel -> ShouldSample)
-> (ThreadLabel -> ShouldSample)
-> ThreadFilter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShouldSample
-> (ThreadLabel -> ShouldSample)
-> Maybe ThreadLabel
-> ShouldSample
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ShouldSample
Yes ((ThreadLabel -> ShouldSample) -> ThreadFilter)
-> (ThreadLabel -> ShouldSample) -> ThreadFilter
forall a b. (a -> b) -> a -> b
$
    Bool -> ShouldSample
fromBool (Bool -> ShouldSample)
-> (ThreadLabel -> Bool) -> ThreadLabel -> ShouldSample
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \ThreadLabel
label ->
      Bool -> Bool
not (Glob
globExclude Glob -> ThreadLabel -> Bool
`matches` ThreadLabel
label)

-- | Construct a thread filter from include and exclude `Glob` patterns.
--
--   If the thread label matches the given include pattern and does not match
--   the given exclude pattern, the thread filter returns `Yes`.
--   Otherwise, the thread filter returns `No`.
--   The thread filter never returns `Never`.
--
--   @since 0.5.0.0
sampleIncludeExclude ::
  -- | The include pattern.
  Glob ->
  -- | The exclude pattern.
  Glob ->
  ThreadFilter
sampleIncludeExclude :: Glob -> Glob -> ThreadFilter
sampleIncludeExclude Glob
globInclude Glob
globExclude =
  (Maybe ThreadLabel -> ShouldSample) -> ThreadFilter
forall a b. a -> b -> a
const ((Maybe ThreadLabel -> ShouldSample) -> ThreadFilter)
-> ((ThreadLabel -> ShouldSample)
    -> Maybe ThreadLabel -> ShouldSample)
-> (ThreadLabel -> ShouldSample)
-> ThreadFilter
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ShouldSample
-> (ThreadLabel -> ShouldSample)
-> Maybe ThreadLabel
-> ShouldSample
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ShouldSample
Yes ((ThreadLabel -> ShouldSample) -> ThreadFilter)
-> (ThreadLabel -> ShouldSample) -> ThreadFilter
forall a b. (a -> b) -> a -> b
$
    Bool -> ShouldSample
fromBool (Bool -> ShouldSample)
-> (ThreadLabel -> Bool) -> ThreadLabel -> ShouldSample
forall b c a. (b -> c) -> (a -> b) -> a -> c
. \ThreadLabel
label ->
      Glob
globInclude Glob -> ThreadLabel -> Bool
`matches` ThreadLabel
label Bool -> Bool -> Bool
&& Bool -> Bool
not (Glob
globExclude Glob -> ThreadLabel -> Bool
`matches` ThreadLabel
label)

-- | Internal helper.
--
--   Construct a `ShouldSample` from a `Bool`.
--
--   Maps `True` to `Yes` and `False` to `No`.
fromBool :: Bool -> ShouldSample
fromBool :: Bool -> ShouldSample
fromBool Bool
b = if Bool
b then ShouldSample
Yes else ShouldSample
No

-- | Read the `Options` from the environment.
--
--   [@GHC_STACK_PROFILER@]:
--     If set to any non-empty value, read and return the options.
--     Otherwise, return `Nothing`, which indicates the `Profiler` should not be started.
--   [@GHC_STACK_PROFILER_WAIT@]:
--     If set to any non-empty value, `wait` is set to `True`.
--   [@GHC_STACK_PROFILER_SAMPLE_INCLUDE@]:
--     If set, `shouldSample` is set to the `ThreadFilter` constructed using `sampleInclude` using the value as a `Glob` pattern.
--     If @GHC_STACK_PROFILER_SAMPLE_EXCLUDE@ is also set, `sampleIncludeExclude` is used.
--   [@GHC_STACK_PROFILER_SAMPLE_EXCLUDE@]:
--     If set, `shouldSample` is set to the `ThreadFilter` constructed using `sampleExclude` using the value as a `Glob` pattern.
--     If @GHC_STACK_PROFILER_SAMPLE_INCLUDE@ is also set, `sampleIncludeExclude` is used.
--   [@GHC_STACK_PROFILER_SAMPLE_RTS_THREADS@]:
--     If set to any non-empty value, `sampleRtsThreads` is set to `True`.
--   [@GHC_STACK_PROFILER_SAMPLE_PROFILER_THREADS@]:
--     If set to any non-empty value, `sampleProfilerThreads` is set to `True`.
--   [@GHC_STACK_PROFILER_SAMPLE_INTERVAL@]:
--     If set to any numeric value, `sampleInterval` is set to the `Interval` constructed using the value as milliseconds.
--     If set to any non-numeric value, a warning is printed to `stderr` and the default `sampleInterval` is used.
--
--   __Warning:__ This function reads environment variables, which is not thread-safe.
--                See [@getenv@](https://en.cppreference.com/c/program/getenv).
--
--   @since 0.5.0.0
fromEnv :: IO (Maybe Options)
fromEnv :: IO (Maybe Options)
fromEnv = do
  shouldStart <- ThreadLabel -> IO Bool
testEnv ThreadLabel
startVar
  if not shouldStart
    then pure Nothing
    else do
      wait <- testEnv waitVar
      shouldSample <-
        (,) <$> lookupEnvGlob sampleIncludeVar <*> lookupEnvGlob sampleExcludeVar <&> \case
          (Maybe Glob
Nothing, Maybe Glob
Nothing) -> Options -> ThreadFilter
shouldSample Options
defaultOptions
          (Just Glob
includeGlob, Maybe Glob
Nothing) -> Glob -> ThreadFilter
sampleInclude Glob
includeGlob
          (Maybe Glob
Nothing, Just Glob
excludeGlob) -> Glob -> ThreadFilter
sampleExclude Glob
excludeGlob
          (Just Glob
includeGlob, Just Glob
excludeGlob) -> Glob -> Glob -> ThreadFilter
sampleIncludeExclude Glob
includeGlob Glob
excludeGlob
      sampleRtsThreads <- testEnv sampleRtsThreadsVar
      sampleProfilerThreads <- testEnv sampleProfilerThreadsVar
      sampleInterval <-
        lookupEnv sampleIntervalVar >>= \case
          Maybe ThreadLabel
Nothing ->
            Interval -> IO Interval
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Interval -> IO Interval) -> Interval -> IO Interval
forall a b. (a -> b) -> a -> b
$ Options -> Interval
sampleInterval Options
defaultOptions
          Just ThreadLabel
sampleIntervalMillisString ->
            case ThreadLabel -> Maybe Int
forall a. Read a => ThreadLabel -> Maybe a
readMaybe ThreadLabel
sampleIntervalMillisString of
              Maybe Int
Nothing -> do
                Handle -> ThreadLabel -> IO ()
hPutStrLn Handle
stderr (ThreadLabel -> IO ()) -> ThreadLabel -> IO ()
forall a b. (a -> b) -> a -> b
$
                  ThreadLabel -> ThreadLabel -> ThreadLabel -> ThreadLabel
forall r. PrintfType r => ThreadLabel -> r
printf
                    ThreadLabel
"Could not parse the value of %s. Expected a number, found %s"
                    ThreadLabel
sampleIntervalVar
                    ThreadLabel
sampleIntervalMillisString
                Interval -> IO Interval
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Interval -> IO Interval) -> Interval -> IO Interval
forall a b. (a -> b) -> a -> b
$ Options -> Interval
sampleInterval Options
defaultOptions
              Just Int
sampleIntervalMillis ->
                Interval -> IO Interval
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (Interval -> IO Interval) -> Interval -> IO Interval
forall a b. (a -> b) -> a -> b
$ Int -> Interval
MkIntervalMillis Int
sampleIntervalMillis
      pure $
        Just
          MkOptions
            { wait
            , shouldSample
            , sampleRtsThreads
            , sampleProfilerThreads
            , sampleInterval
            }
 where
  testEnv :: String -> IO Bool
  testEnv :: ThreadLabel -> IO Bool
testEnv = (Maybe ThreadLabel -> Bool) -> IO (Maybe ThreadLabel) -> IO Bool
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (Bool -> (ThreadLabel -> Bool) -> Maybe ThreadLabel -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (Bool -> Bool
not (Bool -> Bool) -> (ThreadLabel -> Bool) -> ThreadLabel -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ThreadLabel -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null)) (IO (Maybe ThreadLabel) -> IO Bool)
-> (ThreadLabel -> IO (Maybe ThreadLabel))
-> ThreadLabel
-> IO Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ThreadLabel -> IO (Maybe ThreadLabel)
lookupEnv

  lookupEnvGlob :: String -> IO (Maybe Glob)
  lookupEnvGlob :: ThreadLabel -> IO (Maybe Glob)
lookupEnvGlob = (Maybe ThreadLabel -> Maybe Glob)
-> IO (Maybe ThreadLabel) -> IO (Maybe Glob)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ThreadLabel -> Glob) -> Maybe ThreadLabel -> Maybe Glob
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ThreadLabel -> Glob
forall a. IsString a => ThreadLabel -> a
fromString) (IO (Maybe ThreadLabel) -> IO (Maybe Glob))
-> (ThreadLabel -> IO (Maybe ThreadLabel))
-> ThreadLabel
-> IO (Maybe Glob)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. ThreadLabel -> IO (Maybe ThreadLabel)
lookupEnv

  startVar :: String
  startVar :: ThreadLabel
startVar = ThreadLabel
"GHC_STACK_PROFILER"

  waitVar :: String
  waitVar :: ThreadLabel
waitVar = ThreadLabel
"GHC_STACK_PROFILER_WAIT"

  sampleIncludeVar :: String
  sampleIncludeVar :: ThreadLabel
sampleIncludeVar = ThreadLabel
"GHC_STACK_PROFILER_SAMPLE_INCLUDE"

  sampleExcludeVar :: String
  sampleExcludeVar :: ThreadLabel
sampleExcludeVar = ThreadLabel
"GHC_STACK_PROFILER_SAMPLE_EXCLUDE"

  sampleRtsThreadsVar :: String
  sampleRtsThreadsVar :: ThreadLabel
sampleRtsThreadsVar = ThreadLabel
"GHC_STACK_PROFILER_SAMPLE_RTS_THREADS"

  sampleProfilerThreadsVar :: String
  sampleProfilerThreadsVar :: ThreadLabel
sampleProfilerThreadsVar = ThreadLabel
"GHC_STACK_PROFILER_SAMPLE_PROFILER_THREADS"

  sampleIntervalVar :: String
  sampleIntervalVar :: ThreadLabel
sampleIntervalVar = ThreadLabel
"GHC_STACK_PROFILER_SAMPLE_INTERVAL"

-------------------------------------------------------------------------------
-- Low-level API
-------------------------------------------------------------------------------

-------------------------------------------------------------------------------
-- Manager

-- | Run an action with a new `Manager`.
--
--   The first argument indicates if sampler threads should wait for a call to
--  `startProfiling` or a "start profiling" command on the eventlog socket.
--   If you are using @ghc-stack-profiler@ with @eventlog-socket@'s control
--   commands, this should be set to @True@.
--
--   The `Manager` is stopped when the action finishes.
--
--   __Warning:__ This function spawns a `Manager` thread.
--   Having multiple concurrent `Manager` threads is unsupported and unsafe.
--
--   @since 0.5.0.0
withManager ::
  -- | Flag that determines if sampler threads should wait.
  Bool ->
  -- | The action that runs with the `Manager`.
  (Manager -> IO a) ->
  IO a
withManager :: forall a. Bool -> (Manager -> IO a) -> IO a
withManager Bool
wait Manager -> IO a
action =
  IO Manager -> (Manager -> IO ()) -> (Manager -> IO a) -> IO a
forall a b c. IO a -> (a -> IO b) -> (a -> IO c) -> IO c
bracket (Bool -> IO Manager
startManager Bool
wait) Manager -> IO ()
stopManager Manager -> IO a
action

-- | Start a `Manager`.
--
--   The first argument indicates if sampler threads should wait for a call to
--  `startProfiling` or a "start profiling" command on the eventlog socket.
--   If you are using @ghc-stack-profiler@ with @eventlog-socket@'s control
--   commands, this should be set to @True@.
--
--   __Warning:__ This function spawns a `Manager` thread.
--   Having multiple concurrent `Manager` threads is unsupported and unsafe.
--
--   __Warning:__ The manager should be stopped with `stopManager`.
--
--   @since 0.5.0.0
startManager :: Bool -> IO Manager
startManager :: Bool -> IO Manager
startManager Bool
wait = do
  -- TODO: Detect if the event loop thread is running and throw an error.
  manager <- Bool -> IO Manager
newManager Bool
wait
  startEventLoop manager
  Eventlog.Socket.registerWithEventlogSocket manager
  pure manager

-------------------------------------------------------------------------------
-- Sampler
-------------------------------------------------------------------------------

-- | Run an action with a `Sampler` for the current thread.
--
--   The `Sampler` is stopped when the action finishes.
--
--   __Warning:__ If the action creates a new thread, it /will not/ be sampled.
--
--   @since 0.5.0.0
withSamplerForMe :: Manager -> Interval -> (Sampler -> IO a) -> IO a
withSamplerForMe :: forall a. Manager -> Interval -> (Sampler -> IO a) -> IO a
withSamplerForMe Manager
manager Interval
interval Sampler -> IO a
action = do
  IO ThreadId
myThreadId IO ThreadId -> (ThreadId -> IO a) -> IO a
forall a b. IO a -> (a -> IO b) -> IO b
forall (m :: * -> *) a b. Monad m => m a -> (a -> m b) -> m b
>>= \ThreadId
threadId ->
    SamplerDescr -> (Sampler -> IO a) -> IO a
forall a. SamplerDescr -> (Sampler -> IO a) -> IO a
withSampler (Manager -> ThreadId -> Interval -> SamplerDescr
samplerFor Manager
manager ThreadId
threadId Interval
interval) Sampler -> IO a
action

-- | Start a sampler for the given `ThreadId`.
--
--   __Warning:__ The sampler should be stopped using `stopSampler` or `stopManager`.
--
--   @since 0.5.0.0
startSamplerFor :: Manager -> ThreadId -> Interval -> IO Sampler
startSamplerFor :: Manager -> ThreadId -> Interval -> IO Sampler
startSamplerFor Manager
manager ThreadId
threadId Interval
interval =
  SamplerDescr -> IO Sampler
startSampler (Manager -> ThreadId -> Interval -> SamplerDescr
samplerFor Manager
manager ThreadId
threadId Interval
interval)

-- | Internal helper.
--
--   Create a `SamplerDescr` that samples a single thread.
samplerFor :: Manager -> ThreadId -> Interval -> SamplerDescr
samplerFor :: Manager -> ThreadId -> Interval -> SamplerDescr
samplerFor Manager
samplerManager ThreadId
threadId Interval
sampleInterval =
  MkSamplerDescr{Manager
samplerManager :: Manager
samplerManager :: Manager
samplerManager, IO [ThreadId]
samplerThreads :: IO [ThreadId]
samplerThreads :: IO [ThreadId]
samplerThreads, Interval
sampleInterval :: Interval
sampleInterval :: Interval
sampleInterval}
 where
  samplerThreads :: IO [ThreadId]
samplerThreads = [ThreadId] -> IO [ThreadId]
forall a. a -> IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure [ThreadId
threadId]

-- | Start a sampler with the given `Options`.
--
--   This function ignores the `wait` field and uses the value that was
--   passed to the `Manager` on creation.
--
--   __Warning:__ The sampler should be stopped using `stopSampler` or `stopManager`.
--
--   @since 0.5.0.0
startSamplerWith :: Manager -> Options -> IO Sampler
startSamplerWith :: Manager -> Options -> IO Sampler
startSamplerWith Manager
manager Options
options = do
  neverSetRef <- Set ThreadId -> IO (IORef (Set ThreadId))
forall a. a -> IO (IORef a)
newIORef Set ThreadId
forall a. Set a
Set.empty
  startSampler (samplerWith manager neverSetRef options)

-- | Internal helper.
--
--   Create a `SamplerDescr` for the given `Options`.
samplerWith ::
  Manager ->
  IORef (Set ThreadId) ->
  Options ->
  SamplerDescr
samplerWith :: Manager -> IORef (Set ThreadId) -> Options -> SamplerDescr
samplerWith Manager
samplerManager IORef (Set ThreadId)
neverSetRef Options
options =
  MkSamplerDescr{Manager
samplerManager :: Manager
samplerManager :: Manager
samplerManager, IO [ThreadId]
samplerThreads :: IO [ThreadId]
samplerThreads :: IO [ThreadId]
samplerThreads, Interval
sampleInterval :: Interval
sampleInterval :: Interval
sampleInterval}
 where
  MkOptions
    { ThreadFilter
shouldSample :: Options -> ThreadFilter
shouldSample :: ThreadFilter
shouldSample
    , sampleRtsThreads :: Options -> Bool
sampleRtsThreads = Bool -> ShouldSample
fromBool -> ShouldSample
shouldSampleRtsThreads
    , sampleProfilerThreads :: Options -> Bool
sampleProfilerThreads = Bool -> ShouldSample
fromBool -> ShouldSample
shouldSampleProfilerThreads
    , Interval
sampleInterval :: Options -> Interval
sampleInterval :: Interval
sampleInterval
    } = Options
options

  samplerThreads :: IO [ThreadId]
samplerThreads = do
    neverSet <- IORef (Set ThreadId) -> IO (Set ThreadId)
forall a. IORef a -> IO a
readIORef IORef (Set ThreadId)
neverSetRef
    (threadIds', neverSet') <- filterThreads neverSet =<< listThreads
    writeIORef neverSetRef $! neverSet'
    pure threadIds'

  filterThreads :: Set ThreadId -> [ThreadId] -> IO ([ThreadId], Set ThreadId)
  filterThreads :: Set ThreadId -> [ThreadId] -> IO ([ThreadId], Set ThreadId)
filterThreads Set ThreadId
neverSet =
    (([Maybe ThreadId], DList ThreadId) -> ([ThreadId], Set ThreadId))
-> IO ([Maybe ThreadId], DList ThreadId)
-> IO ([ThreadId], Set ThreadId)
forall a b. (a -> b) -> IO a -> IO b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Maybe ThreadId] -> [ThreadId])
-> (DList ThreadId -> Set ThreadId)
-> ([Maybe ThreadId], DList ThreadId)
-> ([ThreadId], Set ThreadId)
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap [Maybe ThreadId] -> [ThreadId]
forall a. [Maybe a] -> [a]
catMaybes ((ThreadId -> Set ThreadId -> Set ThreadId)
-> Set ThreadId -> [ThreadId] -> Set ThreadId
forall a b. (a -> b -> b) -> b -> [a] -> b
forall (t :: * -> *) a b.
Foldable t =>
(a -> b -> b) -> b -> t a -> b
foldr ThreadId -> Set ThreadId -> Set ThreadId
forall a. Ord a => a -> Set a -> Set a
S.insert Set ThreadId
neverSet ([ThreadId] -> Set ThreadId)
-> (DList ThreadId -> [ThreadId]) -> DList ThreadId -> Set ThreadId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. DList ThreadId -> [ThreadId]
DList ThreadId -> [Item (DList ThreadId)]
forall l. IsList l => l -> [Item l]
toList))
      (IO ([Maybe ThreadId], DList ThreadId)
 -> IO ([ThreadId], Set ThreadId))
-> ([ThreadId] -> IO ([Maybe ThreadId], DList ThreadId))
-> [ThreadId]
-> IO ([ThreadId], Set ThreadId)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. WriterT (DList ThreadId) IO [Maybe ThreadId]
-> IO ([Maybe ThreadId], DList ThreadId)
forall w (m :: * -> *) a. Monoid w => WriterT w m a -> m (a, w)
runWriterT
      (WriterT (DList ThreadId) IO [Maybe ThreadId]
 -> IO ([Maybe ThreadId], DList ThreadId))
-> ([ThreadId] -> WriterT (DList ThreadId) IO [Maybe ThreadId])
-> [ThreadId]
-> IO ([Maybe ThreadId], DList ThreadId)
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId))
-> [ThreadId] -> WriterT (DList ThreadId) IO [Maybe ThreadId]
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> [a] -> f [b]
traverse ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
testThread
   where
    testThread :: ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
    testThread :: ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
testThread ThreadId
threadId
      -- If the threadId is in the neverSet, do not sample it.
      | ThreadId
threadId ThreadId -> Set ThreadId -> Bool
forall a. Ord a => a -> Set a -> Bool
`S.member` Set ThreadId
neverSet =
          Maybe ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
forall a. a -> WriterT (DList ThreadId) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ThreadId
forall a. Maybe a
Nothing
      | Bool
otherwise = do
          -- If the threadId is a profiler thread,
          -- it should be sampled if-and-only-if shouldSampleProfilerThreads is true.
          isProfilerThread <- IO Bool -> WriterT (DList ThreadId) IO Bool
forall a. IO a -> WriterT (DList ThreadId) IO a
forall (m :: * -> *) a. MonadIO m => IO a -> m a
liftIO (Manager -> ThreadId -> IO Bool
isProfilerThreadFor Manager
samplerManager ThreadId
threadId)
          if isProfilerThread
            then
              evalShouldSample threadId shouldSampleProfilerThreads
            else do
              maybeThreadLabel <- liftIO (threadLabel threadId)
              -- If the threadId is an RTS thread,
              -- it should be sampled if-and-only-if shouldSampleRtsThreads is true.
              if isRtsThread maybeThreadLabel
                then
                  evalShouldSample threadId shouldSampleRtsThreads
                else
                  -- Otherwise, run the user-provided predicate and follow its instructions.
                  evalShouldSample threadId (shouldSample threadId maybeThreadLabel)

    -- Evaluate a `ShouldSample` judgement for the given threadId.
    evalShouldSample :: ThreadId -> ShouldSample -> WriterT (DList ThreadId) IO (Maybe ThreadId)
    evalShouldSample :: ThreadId
-> ShouldSample -> WriterT (DList ThreadId) IO (Maybe ThreadId)
evalShouldSample ThreadId
threadId = \case
      ShouldSample
Yes -> Maybe ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
forall a. a -> WriterT (DList ThreadId) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure (ThreadId -> Maybe ThreadId
forall a. a -> Maybe a
Just ThreadId
threadId)
      ShouldSample
No -> Maybe ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
forall a. a -> WriterT (DList ThreadId) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ThreadId
forall a. Maybe a
Nothing
      ShouldSample
Never -> DList ThreadId -> WriterT (DList ThreadId) IO ()
forall w (m :: * -> *). (Monoid w, Monad m) => w -> WriterT w m ()
tell ([Item (DList ThreadId)] -> DList ThreadId
forall l. IsList l => [Item l] -> l
fromList [ThreadId
Item (DList ThreadId)
threadId]) WriterT (DList ThreadId) IO ()
-> WriterT (DList ThreadId) IO (Maybe ThreadId)
-> WriterT (DList ThreadId) IO (Maybe ThreadId)
forall a b.
WriterT (DList ThreadId) IO a
-> WriterT (DList ThreadId) IO b -> WriterT (DList ThreadId) IO b
forall (m :: * -> *) a b. Monad m => m a -> m b -> m b
>> Maybe ThreadId -> WriterT (DList ThreadId) IO (Maybe ThreadId)
forall a. a -> WriterT (DList ThreadId) IO a
forall (f :: * -> *) a. Applicative f => a -> f a
pure Maybe ThreadId
forall a. Maybe a
Nothing

-- | Was the given thread created by this library?
isProfilerThreadFor :: Manager -> ThreadId -> IO Bool
isProfilerThreadFor :: Manager -> ThreadId -> IO Bool
isProfilerThreadFor Manager
manager ThreadId
threadId =
  STM Bool -> IO Bool
forall a. STM a -> IO a
atomically (STM Bool -> IO Bool) -> STM Bool -> IO Bool
forall a b. (a -> b) -> a -> b
$ do
    isEventLoopThread <-
      Bool -> Maybe Bool -> Bool
forall a. a -> Maybe a -> a
fromMaybe Bool
False (Maybe Bool -> Bool)
-> (Maybe EventLoop -> Maybe Bool) -> Maybe EventLoop -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (EventLoop -> Bool) -> Maybe EventLoop -> Maybe Bool
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((ThreadId -> ThreadId -> Bool
forall a. Eq a => a -> a -> Bool
== ThreadId
threadId) (ThreadId -> Bool) -> (EventLoop -> ThreadId) -> EventLoop -> Bool
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Async () -> ThreadId
forall a. Async a -> ThreadId
asyncThreadId (Async () -> ThreadId)
-> (EventLoop -> Async ()) -> EventLoop -> ThreadId
forall b c a. (b -> c) -> (a -> b) -> a -> c
. EventLoop -> Async ()
eventLoopAsync)
        (Maybe EventLoop -> Bool) -> STM (Maybe EventLoop) -> STM Bool
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> TVar (Maybe EventLoop) -> STM (Maybe EventLoop)
forall a. TVar a -> STM a
readTVar (Manager -> TVar (Maybe EventLoop)
eventLoopThreadVar Manager
manager)
    isSamplerThread <-
      Map.member threadId
        <$> readTVar (samplerThreadMapVar manager)
    pure $ isEventLoopThread || isSamplerThread

-- | Is the given thread an RTS thread?
isRtsThread :: Maybe ThreadLabel -> Bool
isRtsThread :: Maybe ThreadLabel -> Bool
isRtsThread =
  Bool -> (ThreadLabel -> Bool) -> Maybe ThreadLabel -> Bool
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Bool
False (\ThreadLabel
label -> ThreadLabel
label ThreadLabel -> ThreadLabel -> Bool
forall a. Eq a => a -> a -> Bool
== ThreadLabel
"TimerManager" Bool -> Bool -> Bool
|| ThreadLabel
"IOManager on cap" ThreadLabel -> ThreadLabel -> Bool
forall a. Eq a => [a] -> [a] -> Bool
`isPrefixOf` ThreadLabel
label)