-- | Logging utilities for reporting heap statistics
module Development.IDE.Main.HeapStats ( withHeapStats, Log(..)) where

import           Control.Concurrent
import           Control.Concurrent.Async
import           Control.Monad
import           Data.Word
import           GHC.Stats
import           Ide.Logger               (Pretty (pretty), Priority (Info),
                                           Recorder, WithPriority, hsep,
                                           logWith, (<+>))
import           Text.Printf              (printf)

data Log
  = LogHeapStatsPeriod !Int
  | LogHeapStatsDisabled
  | LogHeapStats !Word64 !Word64
  deriving Int -> Log -> ShowS
[Log] -> ShowS
Log -> String
(Int -> Log -> ShowS)
-> (Log -> String) -> ([Log] -> ShowS) -> Show Log
forall a.
(Int -> a -> ShowS) -> (a -> String) -> ([a] -> ShowS) -> Show a
$cshowsPrec :: Int -> Log -> ShowS
showsPrec :: Int -> Log -> ShowS
$cshow :: Log -> String
show :: Log -> String
$cshowList :: [Log] -> ShowS
showList :: [Log] -> ShowS
Show

instance Pretty Log where
  pretty :: forall ann. Log -> Doc ann
pretty = \case
    LogHeapStatsPeriod Int
period ->
      Doc ann
"Logging heap statistics every" Doc ann -> Doc ann -> Doc ann
forall ann. Doc ann -> Doc ann -> Doc ann
<+> String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Int -> String
toFormattedSeconds Int
period)
    Log
LogHeapStatsDisabled ->
      Doc ann
"Heap statistics are not enabled (RTS option -T is needed)"
    LogHeapStats Word64
liveBytes Word64
heapSize ->
      [Doc ann] -> Doc ann
forall ann. [Doc ann] -> Doc ann
hsep
        [ Doc ann
"Live bytes:"
        , String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Word64 -> String
toFormattedMegabytes Word64
liveBytes)
        , Doc ann
"Heap size:"
        , String -> Doc ann
forall ann. String -> Doc ann
forall a ann. Pretty a => a -> Doc ann
pretty (Word64 -> String
toFormattedMegabytes Word64
heapSize) ]
    where
      toFormattedSeconds :: Int -> String
      toFormattedSeconds :: Int -> String
toFormattedSeconds Int
s = String -> Double -> String
forall r. PrintfType r => String -> r
printf String
"%.2fs" (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Int @Double Int
s Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1e6)

      toFormattedMegabytes :: Word64 -> String
      toFormattedMegabytes :: Word64 -> String
toFormattedMegabytes Word64
b = String -> Double -> String
forall r. PrintfType r => String -> r
printf String
"%.2fMB" (forall a b. (Integral a, Num b) => a -> b
fromIntegral @Word64 @Double Word64
b Double -> Double -> Double
forall a. Fractional a => a -> a -> a
/ Double
1e6)

-- | Interval at which to report the latest heap statistics.
heapStatsInterval :: Int
heapStatsInterval :: Int
heapStatsInterval = Int
60_000_000 -- 60s

-- | Report the live bytes and heap size at the last major collection.
logHeapStats :: Recorder (WithPriority Log) -> IO ()
logHeapStats :: Recorder (WithPriority Log) -> IO ()
logHeapStats Recorder (WithPriority Log)
l = do
  RTSStats
stats <- IO RTSStats
getRTSStats
  -- live_bytes is the total amount of live memory in a program
  -- (corresponding to the amount on a heap profile)
  let live_bytes :: Word64
live_bytes = GCDetails -> Word64
gcdetails_live_bytes (RTSStats -> GCDetails
gc RTSStats
stats)
  -- heap_size is the total amount of memory the RTS is using
  -- this corresponds closer to OS memory usage
      heap_size :: Word64
heap_size  = GCDetails -> Word64
gcdetails_mem_in_use_bytes (RTSStats -> GCDetails
gc RTSStats
stats)
  Recorder (WithPriority Log) -> Priority -> Log -> IO ()
forall (m :: * -> *) msg.
(HasCallStack, MonadIO m) =>
Recorder (WithPriority msg) -> Priority -> msg -> m ()
logWith Recorder (WithPriority Log)
l Priority
Info (Log -> IO ()) -> Log -> IO ()
forall a b. (a -> b) -> a -> b
$ Word64 -> Word64 -> Log
LogHeapStats Word64
live_bytes Word64
heap_size

-- | An action which logs heap statistics at the 'heapStatsInterval'
heapStatsThread :: Recorder (WithPriority Log) -> IO r
heapStatsThread :: forall r. Recorder (WithPriority Log) -> IO r
heapStatsThread Recorder (WithPriority Log)
l = IO () -> IO r
forall (f :: * -> *) a b. Applicative f => f a -> f b
forever (IO () -> IO r) -> IO () -> IO r
forall a b. (a -> b) -> a -> b
$ do
  Int -> IO ()
threadDelay Int
heapStatsInterval
  Recorder (WithPriority Log) -> IO ()
logHeapStats Recorder (WithPriority Log)
l

-- | A helper function which launches the 'heapStatsThread' and kills it
-- appropriately when the inner action finishes. It also checks to see
-- if `-T` is enabled.
withHeapStats :: Recorder (WithPriority Log) -> IO r -> IO r
withHeapStats :: forall r. Recorder (WithPriority Log) -> IO r -> IO r
withHeapStats Recorder (WithPriority Log)
l IO r
k = do
  Bool
enabled <- IO Bool
getRTSStatsEnabled
  if Bool
enabled
    then do
      Recorder (WithPriority Log) -> Priority -> Log -> IO ()
forall (m :: * -> *) msg.
(HasCallStack, MonadIO m) =>
Recorder (WithPriority msg) -> Priority -> msg -> m ()
logWith Recorder (WithPriority Log)
l Priority
Info (Log -> IO ()) -> Log -> IO ()
forall a b. (a -> b) -> a -> b
$ Int -> Log
LogHeapStatsPeriod Int
heapStatsInterval
      IO Any -> (Async Any -> IO r) -> IO r
forall a b. IO a -> (Async a -> IO b) -> IO b
withAsync (Recorder (WithPriority Log) -> IO Any
forall r. Recorder (WithPriority Log) -> IO r
heapStatsThread Recorder (WithPriority Log)
l) (IO r -> Async Any -> IO r
forall a b. a -> b -> a
const IO r
k)
    else do
      Recorder (WithPriority Log) -> Priority -> Log -> IO ()
forall (m :: * -> *) msg.
(HasCallStack, MonadIO m) =>
Recorder (WithPriority msg) -> Priority -> msg -> m ()
logWith Recorder (WithPriority Log)
l Priority
Info Log
LogHeapStatsDisabled
      IO r
k