ThreadSample.hs

{-# LANGUAGE OverloadedStrings #-}

module GHC.Stack.Profiler.Core.ThreadSample (
  -- * High-level API
  ThreadSample (..),
  deserializeEventlogMessage,

  -- * Serialisable 'ThreadSample'
  CallStackMessage (..),
  StackItem (..),
  SourceLocation (..),

  -- * Serialisation of 'CallStackMessage'
  SymbolTableWriter (..),
  SymbolTableReader (..),
  dehydrateCallStackMessage,
  BinaryCallStackDecodeError (..),
  hydrateEventlogCallStackMessage,
  catCallStackMessage,
  chunkCallStackMessage,
  chunkCallStackMessage_,

  -- * Message dehydration helpers
  EncodingState,
  runWithEncodingState,
  newEncodingState,
  lookupSourceLocationMessage,
  lookupTextMessage,
) where

import Control.Concurrent
import Control.Exception (Exception (..))
import Control.Monad (when)
import Control.Monad.Trans.State.Strict (State, runState)
import qualified Control.Monad.Trans.State.Strict as State
import Data.Binary
import Data.Binary.Get
import qualified Data.ByteString.Lazy as LBS
import Data.Either (partitionEithers)
import qualified Data.List as List
import Data.List.NonEmpty (NonEmpty (..))
import qualified Data.List.NonEmpty as NonEmpty
import Data.Text (Text)
import qualified Data.Text as Text
import GHC.Generics

import GHC.Stack.CloneStack (StackSnapshot)
import GHC.Stack.Profiler.Core.Eventlog
import GHC.Stack.Profiler.Core.SourceLocation
import GHC.Stack.Profiler.Core.SymbolTable
import GHC.Stack.Profiler.Core.Util (word16ToInt)

-- ----------------------------------------------------------------------------
-- Thread Sample
-- ----------------------------------------------------------------------------

-- | A 'ThreadSample' is a snapshot of a threads RTS callstack.
-- This callstack is a copy of the original callstack, so can be traversed and
-- decoded without affecting the running thread.
--
-- The 'StackSnapshot' is a boxed value and needs to be garbage collected.
-- Note, as long as 'StackSnapshot' is alive, you keep the full callstack
-- alive, which might be quite expensive.
data ThreadSample = ThreadSample
  { threadSampleId :: !ThreadId
  , threadSampleCapability :: !CapabilityId
  , threadSampleStackSnapshot :: !StackSnapshot
  }
  deriving (Generic)

deserializeEventlogMessage :: LBS.ByteString -> Either String BinaryEventlogMessage
deserializeEventlogMessage msg = case runGetOrFail get msg of
  Left (_, _, errMsg) -> Left errMsg
  Right (_, _, callStackMessage) -> Right callStackMessage

-- ----------------------------------------------------------------------------
-- Decoded RTS CallStack
-- ----------------------------------------------------------------------------

-- | A decoded rts callstack that can be serialised to the EventLog.
data CallStackMessage = MkCallStackMessage
  { callThreadId :: !Word64
  , callCapabilityId :: !CapabilityId
  , callStack :: [StackItem]
  }
  deriving (Eq, Ord, Show, Generic)

data StackItem
  = IpeId !IpeId
  | UserAnnotation !String !(Maybe SourceLocation)
  deriving (Eq, Ord, Show, Generic)

-- ----------------------------------------------------------------------------
-- Turning a 'CallStackMessage' into '[BinaryEventlogMessage]'
-- ----------------------------------------------------------------------------

-- | Generic implementation to turn 'CallStackMessage' into '[BinaryEventlogMessage]'.
--
-- Replaces already encountered text or source location information with unique ids.
-- If new text or source location messages are encountered, they are inserted into
-- the 'SymbolTableWriter'.
--
-- All new string values and source location messages are before 'CallStackChunk' and
-- 'CallStackFinal' messages.
-- For the result list @r :: ['BinaryEventlogMessage']@, the following holds:
--
-- * 'StringDef' messages are the first elements in @r@. There might not be any.
-- * 'SourceLocationDef' are after 'StringDef' messages and before any 'CallStackChunk' or
--    'CallStackFinal' messages. There might not be any such messages.
-- * Then 'CallStackChunk' follow if there are any.
-- * The last message is always a 'CallStackFinal' message and it occurs exactly once in @r@.
dehydrateCallStackMessage ::
  forall table.
  SymbolTableWriter table ->
  CallStackMessage ->
  ([BinaryEventlogMessage], SymbolTableWriter table)
dehydrateCallStackMessage msgTbl0 msg =
  let
    (stackItems, finalState) =
      runWithEncodingState
        (newEncodingState msgTbl0)
        (mapM go (callStack msg))

    stringDefs =
      map StringDef $ stringMessages finalState

    sourceLocDefs =
      map SourceLocationDef $ sourceLocMessages finalState

    -- TODO: this needs to be fixed
    stackMsgChunks =
      chunkCallStackMessage
        MkBinaryCallStackMessage
          { binaryCallThreadId = callThreadId msg
          , binaryCallCapabilityId = callCapabilityId msg
          , binaryCallStack = stackItems
          }
  in
    ( stringDefs ++ sourceLocDefs ++ stackMsgChunks
    , symbolTableWriter finalState
    )
 where
  go :: StackItem -> State (EncodingState tbl) BinaryStackItem
  go = \case
    IpeId ipeId ->
      pure $ BinaryIpe ipeId
    UserAnnotation s mSrcLoc -> do
      srcLocId <- case mSrcLoc of
        Nothing -> pure Nothing
        Just srcLoc -> Just <$> lookupSourceLocationMessage srcLoc
      BinaryMessage <$> lookupTextMessage (Text.pack s) <*> pure srcLocId

data BinaryCallStackDecodeError
  = StringIdNotFound StringId
  | SourceLocationIdNotFound SourceLocationId
  deriving (Show)

instance Exception BinaryCallStackDecodeError where
  displayException = \case
    StringIdNotFound sid ->
      "Failed to decode a BinaryCallStackMessage. Failed to find a String with the key: " ++ show (getStringId sid)
    SourceLocationIdNotFound sid ->
      "Failed to decode a BinaryCallStackMessage. Failed to find a SourceLocation with the key: " ++ show (getSourceLocationId sid)

-- | Generic implementation to turn 'BinaryCallStackMessage' into the much richer
-- 'CallStackMessage'.
hydrateEventlogCallStackMessage :: SymbolTableReader -> BinaryCallStackMessage -> (CallStackMessage, [BinaryCallStackDecodeError])
hydrateEventlogCallStackMessage decodeTable msg =
  let
    decodeItem :: BinaryStackItem -> Either BinaryCallStackDecodeError StackItem
    decodeItem = \case
      BinaryIpe ipeId ->
        Right $ IpeId ipeId
      BinaryMessage stringId mSrcLocId -> do
        str <-
          maybe
            (Left $ StringIdNotFound stringId)
            (Right . Text.unpack)
            (lookupStringId decodeTable stringId)
        srcLoc <- case mSrcLocId of
          Nothing -> pure Nothing
          Just srcLocId ->
            maybe
              (Left $ SourceLocationIdNotFound srcLocId)
              (Right . Just)
              (lookupSourceLocationId decodeTable srcLocId)
        pure $ UserAnnotation str srcLoc

    itemsOrErros = map decodeItem (binaryCallStack msg)
    (errors, items) = partitionEithers itemsOrErros
  in
    ( MkCallStackMessage
        { callCapabilityId = binaryCallCapabilityId msg
        , callThreadId = binaryCallThreadId msg
        , callStack = items
        }
    , errors
    )

-- | Combine all 'BinaryCallStackMessage's into a single 'BinaryCallStackMessage'.
-- We assume that all 'BinaryCallStackMessage' only differ in their 'binaryCallStack' values.
--
-- 'catCallStackMessage' is the conceptually inverse of 'chunkCallStackMessage'.
catCallStackMessage :: NonEmpty BinaryCallStackMessage -> BinaryCallStackMessage
catCallStackMessage msgs =
  MkBinaryCallStackMessage
    { binaryCallThreadId = binaryCallThreadId $ NonEmpty.head msgs
    , binaryCallCapabilityId = binaryCallCapabilityId $ NonEmpty.head msgs
    , binaryCallStack = concatMap (reverse . binaryCallStack) . reverse $ NonEmpty.toList msgs
    }

-- | Chunk the 'binaryCallStack' of the 'BinaryCallStackMessage' by the given 'Word16'.
-- If there are no items in 'BinaryCallStackMessage', then a singleton list is returned containing
-- the original element.
--
-- Post-condition for the result @r@:
--
-- * all elements in @init r @ are 'CallStackChunk's
-- * the element returned by @last r@ is a 'CallStackFinal' BinaryEventlogMessage.
--
-- The resulting 'CallStackChunk' are in reverse order and so are the chunks themselves.
--
-- This means, for a stack @[1,2,3,4,5,6]@ and an assumed chunk size of 2,
-- we produce @[[6,5],[4,3],[2,1]]@.
chunkCallStackMessage :: BinaryCallStackMessage -> [BinaryEventlogMessage]
chunkCallStackMessage = chunkCallStackMessage_ callStackSizeLimit

-- | Same as 'chunkCallStackMessage', but allows to set the chunking size in bytes.
chunkCallStackMessage_ :: Word16 -> BinaryCallStackMessage -> [BinaryEventlogMessage]
chunkCallStackMessage_ chunkLimit16 msg0 =
  let
    chunkLimitInt = word16ToInt chunkLimit16
    items = binaryCallStack msg0
    chunked =
      let
        go (!size, curChunk, restChunk) item =
          let
            !bytes = word16ToInt $ byteSizeOf item
          in
            if (size + bytes) < chunkLimitInt
              then (size + bytes, item : curChunk, restChunk)
              else (bytes, [item], curChunk : restChunk)
        (_, lastChunk, initChunk) = List.foldl' go (0, [], []) items
      in
        lastChunk : initChunk
  in
    mkEventlogMessages chunked
 where
  mkCallStack chunk =
    MkBinaryCallStackMessage
      { binaryCallThreadId = binaryCallThreadId msg0
      , binaryCallCapabilityId = binaryCallCapabilityId msg0
      , binaryCallStack = chunk
      }

  mkEventlogMessages :: [[BinaryStackItem]] -> [BinaryEventlogMessage]
  mkEventlogMessages [] =
    -- If there are no chunks, we simply return the original message
    [ CallStackFinal msg0
    ]
  mkEventlogMessages [chunk] =
    [ CallStackFinal $ mkCallStack chunk
    ]
  mkEventlogMessages (chunk : chunks) =
    CallStackChunk (mkCallStack chunk) : mkEventlogMessages chunks

-- ----------------------------------------------------------------------------
-- Helper types and functions to implement the conversion to the binary
-- representation.
-- ----------------------------------------------------------------------------

data EncodingState tbl = MkEncodingState
  { symbolTableWriter :: !(SymbolTableWriter tbl)
  , stringMessages :: ![BinaryStringMessage]
  , sourceLocMessages :: ![BinarySourceLocationMessage]
  }
  deriving (Generic)

runWithEncodingState :: EncodingState tbl -> State (EncodingState tbl) a -> (a, EncodingState tbl)
runWithEncodingState encodingState encoder =
  runState encoder encodingState

newEncodingState :: SymbolTableWriter tbl -> EncodingState tbl
newEncodingState msgTbl0 =
  MkEncodingState
    { symbolTableWriter = msgTbl0
    , stringMessages = []
    , sourceLocMessages = []
    }

setSymbolTableWriter :: tbl -> State.State (EncodingState tbl) ()
setSymbolTableWriter tbl = State.modify' (\st -> st{symbolTableWriter = (symbolTableWriter st){writerTable = tbl}})

addStringMessage :: BinaryStringMessage -> State.State (EncodingState tbl) ()
addStringMessage msg = State.modify' (\st -> st{stringMessages = msg : stringMessages st})

addSourceLocationMessage :: BinarySourceLocationMessage -> State.State (EncodingState tbl) ()
addSourceLocationMessage msg = State.modify' (\st -> st{sourceLocMessages = msg : sourceLocMessages st})

lookupOrInsertTextMessage :: forall tbl. Text -> State (EncodingState tbl) (StringId, Bool)
lookupOrInsertTextMessage s = do
  tbl <- State.gets symbolTableWriter
  let
    (sid, new, tbl1) = lookupOrInsertText tbl (writerTable tbl) s
  setSymbolTableWriter tbl1
  pure (sid, new)

lookupOrInsertSrcLocMessage :: forall tbl. SourceLocation -> State (EncodingState tbl) (SourceLocationId, Bool)
lookupOrInsertSrcLocMessage s = do
  tbl <- State.gets symbolTableWriter
  let
    (sid, new, tbl1) = lookupOrInsertSourceLocation tbl (writerTable tbl) s
  setSymbolTableWriter tbl1
  pure (sid, new)

lookupTextMessage :: forall tbl. Text -> State (EncodingState tbl) StringId
lookupTextMessage s = do
  (sid, new) <- lookupOrInsertTextMessage s
  when new $
    addStringMessage $
      MkBinaryStringMessage sid s
  pure sid

lookupSourceLocationMessage :: forall tbl. SourceLocation -> State (EncodingState tbl) SourceLocationId
lookupSourceLocationMessage s = do
  (sid, new) <- lookupOrInsertSrcLocMessage s
  when new $ do
    fileId <- lookupTextMessage $ fileName s
    addSourceLocationMessage $
      MkBinarySourceLocationMessage
        { binarySourceLocationMessageId = sid
        , binarySourceLocationRow = line s
        , binarySourceLocationColumn = column s
        , binarySourceLocationFilename = fileId
        }
  pure sid