ParTrace.hs

{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE NoApplicativeDo #-}
-- | Functions to support the constant space traversal of a heap.
-- This module is like the Trace module but performs the tracing in
-- parellel. The speed-up is quite modest but hopefully can be improved in
-- future.
--
-- The tracing functions create a thread for each MBlock which we
-- traverse, closures are then sent to the relevant threads to be
-- dereferenced and thread-local storage is accumulated.
module GHC.Debug.ParTrace ( traceParFromM
                          , tracePar
                          , traceParFromWithState
                          , asyncTraceParFromWithState
                          , justClosuresPar
                          , WorkerId(..)
                          , TraceFunctionsIO(..)
                          , ClosurePtrWithInfo(..)
                          ) where

import GHC.Debug.Async
import GHC.Debug.Client.Monad.Class
import GHC.Debug.Client.Monad.Simple
import GHC.Debug.Client.Query
import GHC.Debug.Types

import Control.Concurrent.Async
import Control.Concurrent.STM
import Control.Exception.Base
import Control.Monad
import Control.Monad.Reader
import Data.Array.BitArray.IO hiding (map)
import Data.Bitraversable
import Data.Coerce (coerce)
import Data.IORef
import qualified Data.IntMap as IM
import qualified Data.Map as Map
import Data.Word
import GHC.Conc (numCapabilities)

threads :: Int
threads = numCapabilities

type InChan = TChan
type OutChan = TChan

-- | State local to a thread, there are $threads spawned, each which deals
-- with (address `div` 8) % threads. Each thread therefore:
--
-- * Outer map, segmented by MBlock
--  * Inner map, blocks for that MBlock
--    * Inner IOBitArray, visited information for that block
data ThreadState = ThreadState
  { visitedPtrs :: IM.IntMap (IM.IntMap (IOBitArray Word16))
  }

newtype ThreadInfo a = ThreadInfo (InChan (ClosurePtrWithInfo a))

-- | A 'ClosurePtr' with some additional information which needs to be
-- communicated across to another thread.
data ClosurePtrWithInfo a = ClosurePtrWithInfo !a !ClosurePtr
                          | CCSPtrWithInfo CCSPtr

-- | Map from Thread -> Information about the thread
type ThreadMap a = IM.IntMap (ThreadInfo a)

newtype TraceState a = TraceState { visited :: (ThreadMap a) }


getKeyTriple :: ClosurePtr -> (Int, Int, Word16)
getKeyTriple cp =
  let BlockPtr raw_bk = applyBlockMask cp
      bk = fromIntegral raw_bk `div` 8
      offset = getBlockOffset cp `div` 8
      BlockPtr raw_mbk = applyMBlockMask cp
      mbk = fromIntegral raw_mbk `div` 8
  in (mbk, bk, fromIntegral offset)

getMBlockKey :: ClosurePtr -> WorkerId
getMBlockKey cp =
  let BlockPtr raw_bk = applyMBlockMask cp
  -- Not sure why I had to divide this by 4, but I did.
  in WorkerId ((fromIntegral raw_bk `div` fromIntegral mblockMask `div` 4) `mod` threads)

sendToChan :: TraceState a -> ClosurePtrWithInfo a -> DebugM ()
sendToChan  ts cpi@(ClosurePtrWithInfo _ cp) = DebugM $ liftIO $ do
  let st = visited ts
      WorkerId mkey = getMBlockKey cp
  case IM.lookup mkey st of
    Nothing -> error $ "Not enough chans:" ++ show mkey ++ show threads
    Just (ThreadInfo ic) -> atomically $ writeTChan ic cpi
sendToChan ts cpi@(CCSPtrWithInfo ccsPtr) = DebugM $ liftIO $ do
  let st = visited ts
      WorkerId mkey = getMBlockKey (coerce ccsPtr)
  case IM.lookup mkey st of
    Nothing -> error $ "Not enough chans:" ++ show mkey ++ show threads
    Just (ThreadInfo ic) -> atomically $ writeTChan ic cpi

initThread :: WorkerId
           -> TraceFunctionsIO a ()
           -> DebugM (ThreadInfo a, STM Bool, (ClosurePtrWithInfo a -> DebugM ()) -> IO ())
initThread n k = DebugM $ do
  e <- ask
  ic <- liftIO $ newTChanIO
  let oc = ic
  worker_active <- liftIO $ newTVarIO True
  let start go = runSimple e $ workerThread n k worker_active go oc
      finished = do
        active <- not <$> readTVar worker_active
        empty  <- isEmptyTChan ic
        return (active && empty)

  return (ThreadInfo ic, finished, start)

workerThread :: forall a . WorkerId -> TraceFunctionsIO a () -> TVar Bool -> (ClosurePtrWithInfo a -> DebugM ()) -> OutChan (ClosurePtrWithInfo a) -> DebugM ()
workerThread n k worker_active go oc = DebugM $ do
  d <- ask
  r <- liftIO $ newIORef $
    ThreadState
      { visitedPtrs = IM.empty
      }

  liftIO $ runSimple d (loop r)
  where
    loop r = do
      mcp <- unsafeLiftIO $ try $ do
              atomically $ writeTVar worker_active False
              atomically $ do
                v <- readTChan oc
                writeTVar worker_active True
                return v
      case mcp of
        -- The thread gets blocked on readChan when the work is finished so
        -- when this happens, catch the exception and return the accumulated
        -- state for the thread. Each thread has a reference to all over
        -- threads, so the exception is only raised when ALL threads are
        -- waiting for work.
        Left AsyncCancelled -> return ()
        Right cpi -> deref r cpi >> loop r

    deref r (ClosurePtrWithInfo a cp) = do
        m <- unsafeLiftIO $ readIORef r
        do
          (m', b) <- unsafeLiftIO $ checkVisit cp m
          unsafeLiftIO $ writeIORef r m'
          if b
            then do
              visitedClosVal k n cp a
            else do
              sc <- dereferenceClosure cp
              (a', (), cont) <- closTrace k n cp sc a
              cont (() <$ hextraverse (goCCS r) (gosrt r a') (gop r a') gocd (gos r a') (goc r . ClosurePtrWithInfo a') sc)
    deref r (CCSPtrWithInfo cp) = do
        m <- unsafeLiftIO $ readIORef r
        do
          (m', b) <- unsafeLiftIO $ checkVisit (coerce cp) m
          unsafeLiftIO $ writeIORef r m'
          if b
            then do
              visitedCcsVal k n cp
            else do
              ccs' <- dereferenceCCS cp
              ccsTrace k n cp ccs'
              () <$ bitraverse (goCCS r) goCC ccs'

    goc r c@(ClosurePtrWithInfo _i cp) =
      let mkey = getMBlockKey cp
      in if (mkey == n)
          then deref r c
          else go c

    -- Just do the other dereferencing in the same thread for other closure
    -- types as they are not as common.
    gos r a st = do
      st' <- dereferenceStack st
      stackTrace k n st'
      () <$ bitraverse (gosrt r a) (goc r . ClosurePtrWithInfo a) st'

    gocd d = do
      cd <- dereferenceConDesc d
      conDescTrace k n cd

    goCCS r cp =
        let mkey = getMBlockKey (coerce cp)
        in if (mkey == n)
            then deref r (CCSPtrWithInfo cp)
            else go (CCSPtrWithInfo cp)

    goCC p = do
      () <$ dereferenceCC p

    gop r a p = do
      p' <- dereferencePapPayload p
      papTrace k n p'
      () <$ traverse (goc r . ClosurePtrWithInfo a) p'

    gosrt r a p = do
      p' <- dereferenceSRT p
      srtTrace k n p'
      () <$ traverse (goc r . ClosurePtrWithInfo a) p'


handleBlockLevel :: IM.Key
                    -> Word16
                    -> IM.IntMap (IOBitArray Word16)
                    -> IO (IM.IntMap (IOBitArray Word16), Bool)

handleBlockLevel bk offset m = do
  case IM.lookup bk m of
    Nothing -> do
      na <- newArray (0, fromIntegral (blockMask `div` 8)) False
      writeArray na offset True
      return (IM.insert bk na m, False)
    Just bm -> do
      res <- readArray bm offset
      unless res (writeArray bm offset True)
      return (m, res)

checkVisit :: ClosurePtr -> ThreadState -> IO (ThreadState, Bool)
checkVisit cp st = do
  let (mbk, bk, offset) = getKeyTriple cp
      ThreadState v = st
  case IM.lookup mbk v of
    Nothing -> do
      (st', res) <- handleBlockLevel bk offset IM.empty
      return (ThreadState (IM.insert mbk st' v) , res)
    Just bm -> do
      (st', res) <- handleBlockLevel bk offset bm
      return (ThreadState (IM.insert mbk st' v) , res)


newtype WorkerId = WorkerId Int deriving (Eq, Ord)

data TraceFunctionsIO a s =
      TraceFunctionsIO { papTrace :: !(WorkerId -> GenPapPayload ClosurePtr -> DebugM ())
      , srtTrace :: !(WorkerId -> GenSrtPayload ClosurePtr -> DebugM ())
      , stackTrace :: !(WorkerId -> GenStackFrames SrtCont ClosurePtr -> DebugM ())
      , closTrace :: !(WorkerId -> ClosurePtr -> SizedClosure -> a -> DebugM (a, s, DebugM () -> DebugM ()))
      , visitedClosVal :: !(WorkerId -> ClosurePtr -> a -> DebugM s)
      , visitedCcsVal :: !(WorkerId -> CCSPtr -> DebugM s)
      , conDescTrace :: !(WorkerId -> ConstrDesc -> DebugM ())
      , ccsTrace :: !(WorkerId -> CCSPtr -> CCSPayload -> DebugM s)
      }

justClosuresPar :: Monoid s =>
                   (WorkerId -> ClosurePtr -> SizedClosure -> a -> DebugM (a, s, DebugM () -> DebugM ()))
                -> TraceFunctionsIO a s
justClosuresPar clos = TraceFunctionsIO
  { papTrace = nop
  , srtTrace = nop
  , stackTrace = nop
  , closTrace = clos
  , visitedClosVal = const (const (const (return mempty)))
  , visitedCcsVal = const (const (return mempty))
  , conDescTrace = nop
  , ccsTrace = const (const (const (return mempty)))
  }
  where
    nop = const (const (return ()))

-- | Accumulate state from each node
stateAccumulator  :: Semigroup s => Map.Map WorkerId (IORef s)
                  -> TraceFunctionsIO a s
                  -> TraceFunctionsIO a ()
stateAccumulator refs f =
    TraceFunctionsIO
      { papTrace = papTrace f
      , srtTrace = srtTrace f
      , stackTrace = stackTrace f
      , closTrace = \wid cp sc a -> do
          let ref = refs Map.! wid
          closTrace f wid cp sc a >>= \ (a', s, cont) -> do
            unsafeLiftIO $ modifyIORef' ref (s <>)
            return (a', (), cont)


      , visitedClosVal = \wid cp a -> do
          let ref = refs Map.! wid
          visitedClosVal f wid cp a >>= \ s -> do
            unsafeLiftIO $ modifyIORef' ref (s <>)
            return ()
      , visitedCcsVal = \wid ccs -> do
          let ref = refs Map.! wid
          visitedCcsVal f wid ccs >>= \ s -> do
            unsafeLiftIO $ modifyIORef' ref (s <>)
            return ()
      , conDescTrace = conDescTrace f
      , ccsTrace = \wid ccs ccs' -> do
          let ref = refs Map.! wid
          ccsTrace f wid ccs ccs' >>= \ s -> do
            unsafeLiftIO $ modifyIORef' ref (s <>)
            return ()
      }

-- | Perform a parallel trace, accumulating state as we go, which can be queried whilst
-- the traversal is running.
asyncTraceParFromWithState :: Monoid s
                           => TraceFunctionsIO a s
                           -- ^ The tracing function for dealing with closures.
                           -> [ClosurePtrWithInfo a]
                           -> DebugM (AsyncTrace s)
                           -- An IO action which can be used to query the current result of the trace, and an action to run the trace.
asyncTraceParFromWithState f cps = do
  refs <- Map.fromList <$> mapM (\i -> (WorkerId i,) <$> unsafeLiftIO (newIORef mempty)) [0 .. threads - 1]
  -- Collect the results
  let read_results = mconcat <$> mapM (\(_wid, ref) -> readIORef ref) (Map.toList refs)
  withAsyncTrace read_results (traceParFromM (stateAccumulator refs f) cps)

-- | A parallel traversal which accumulates state when visiting closures.
traceParFromWithState :: Monoid s => TraceFunctionsIO a s
                        -> [ClosurePtrWithInfo a]
                        -> DebugM s
traceParFromWithState f cps = do
  asyncTrace <- asyncTraceParFromWithState f cps
  waitForAsyncTraceResultM asyncTrace

-- | A generic heap traversal function which will use a small amount of
-- memory linear in the heap size. Using this function with appropriate
-- accumulation functions you should be able to traverse quite big heaps in
-- not a huge amount of memory.
--
-- The performance of this parallel version depends on how much contention
-- the functions given in 'TraceFunctionsIO' content for the handle
-- connecting for the debuggee (which is protected by an 'MVar'). With no
-- contention, and precached blocks, the workload can be very evenly
-- distributed leading to high core utilisation.
--
-- As performance depends highly on contention, snapshot mode is much more
-- amenable to parallelisation where the time taken for requests is much
-- lower.
traceParFromM :: TraceFunctionsIO a () -> [ClosurePtrWithInfo a] -> DebugM ()
traceParFromM k cps = do
  traceMsg ("SPAWNING: " ++ show threads)
  -- Set up the worker threads.
  let workerThreadIds = [0 .. threads - 1]
      setupWorker wid = do
        (ti, working, start) <- initThread (WorkerId wid) k
        return ((fromIntegral wid, ti), working, start)
  (init_mblocks, work_actives, start) <- unzip3 <$> mapM setupWorker workerThreadIds
  let ts_map = IM.fromList init_mblocks
      go = sendToChan (TraceState ts_map)

  -- Start the work!
  -- First, we start populating the main queue with tasks.
  -- Then, we start the worker threads concurrenly, which ensures they all get cancelled
  -- if one of the worker threads crash in any way.
  -- At last, we wait for completion in a coordinated fashion.
  dbg <- DebugM ask
  unsafeLiftIO $
    withAsync (mapConcurrently_ ($ go) start) $ \ workerPool -> do
      runSimple dbg $ mapM_ go cps
      -- We have to wait at least for
      withAsync (waitFinish work_actives) $ \ allWorkersFinished -> do
        result <- waitEitherCancel allWorkersFinished workerPool
        case result of
          Left () -> pure ()
          Right () -> throwIO $ ErrorCall $ "traceParFromM: Worker threads terminated"

waitFinish :: [STM Bool] -> IO ()
waitFinish working = atomically (checkDone working)
  where
    checkDone [] = return ()
    checkDone (x:xs) = do
      b <- x
      -- The variable tracks whether the thread thinks it's finished (no
      -- active work and empty chan)
      if b then checkDone xs else retry

-- | A parallel tracing function.
tracePar :: [ClosurePtr] -> DebugM ()
tracePar = traceParFromM funcs . map (ClosurePtrWithInfo ())
  where
    nop = const (const (return ()))
    funcs = TraceFunctionsIO
      { papTrace = nop
      , srtTrace = nop
      , stackTrace = stack
      , closTrace = clos
      , visitedClosVal = const (const (const (return ())))
      , visitedCcsVal = nop
      , conDescTrace = nop
      , ccsTrace = const (const (const (return ())))
      }

    stack :: WorkerId -> GenStackFrames SrtCont ClosurePtr -> DebugM ()
    stack _ fs =
      let stack_frames = getFrames fs
      in mapM_ (getSourceInfo . tableId . frame_info) stack_frames

    clos :: WorkerId -> ClosurePtr -> SizedClosure -> ()
              -> DebugM ((),  (), DebugM () -> DebugM ())
    clos _ _cp sc _ = do
      let itb = info (noSize sc)
      _traced <- getSourceInfo (tableId itb)
      return ((), (), id)