Merge pull request #5883 from unisonweb/topic/thread-profile

Author: aryairani

parser-typechecker/src/Unison/Codebase/Runtime/Profile.hs

@@ -37,9 +37,24 @@ import Unison.Util.Pretty as P
 newtype ProfTrie k a = ProfT (Map k (a, ProfTrie k a))
   deriving (Functor)
 
+trimEmpty :: (Ord k, Eq a, Num a) => ProfTrie k a -> ProfTrie k a
+trimEmpty (ProfT m) = case M.traverseMaybeWithKey f m of
+  Identity m -> ProfT m
+  where
+    f _ (a, trimEmpty -> sub@(ProfT m))
+      | a == 0, null m = pure Nothing
+      | otherwise = pure $ Just (a, sub)
+
+demux ::
+  (Ord k, Eq a, Eq b, Num a, Num b) =>
+  ProfTrie k (a, b) ->
+  (ProfTrie k a, ProfTrie k b)
+demux tr = (trimEmpty $ fst <$> tr, trimEmpty $ snd <$> tr)
+
 -- A profile pairs the above arbitrary key based profile trie with a
 -- decoding of the integers to references and a total sample count.
-data Profile k = Prof !Int !(ProfTrie k Int) !(Map k Reference)
+data Profile k
+  = Prof !(Int, Int) !(ProfTrie k (Int, Int)) !(Map k Reference)
 
 -- Abstracts over the exact key type used in a profile.
 data SomeProfile = forall k. (Ord k) => SomeProf (Profile k)
@@ -48,28 +63,51 @@ data ProfileSpec = NoProf | MiniProf | FullProf String
   deriving (Eq, Ord, Show)
 
 emptyProfile :: Profile k
-emptyProfile = Prof 0 (ProfT M.empty) M.empty
+emptyProfile = Prof zero (ProfT M.empty) M.empty
 
--- Creates a singleton profile trie from a path.
-singlePath :: (Ord k) => [k] -> (Int, ProfTrie k Int)
-singlePath [] = (1, ProfT M.empty)
-singlePath (i : is) = (0,) . ProfT $! M.singleton i (singlePath is)
+zero :: (Int, Int)
+zero = (0, 0)
+
+inc :: Bool -> (Int, Int) -> (Int, Int)
+inc b (m, n) = pair (m + 1) (if b then n + 1 else n)
+  where
+    pair !x !y = (x, y)
 
-addPath0 :: (Ord k) => [k] -> (Int, ProfTrie k Int) -> (Int, ProfTrie k Int)
-addPath0 [] (m, p) = (,p) $! m + 1
-addPath0 (i : is) (m, ProfT p) = (m,) . ProfT $! M.alter f i p
+-- Creates a singleton profile trie from a path.
+singlePath ::
+  (Ord k) =>
+  Bool ->
+  [k] ->
+  ((Int, Int), ProfTrie k (Int, Int))
+singlePath b [] = (inc b zero, ProfT M.empty)
+singlePath b (i : is) =
+  (zero,) . ProfT $! M.singleton i (singlePath b is)
+
+addPath0 ::
+  (Ord k) =>
+  Bool ->
+  [k] ->
+  ((Int, Int), ProfTrie k (Int, Int)) ->
+  ((Int, Int), ProfTrie k (Int, Int))
+addPath0 b [] (t, p) = (,p) $! inc b t
+addPath0 b (i : is) (m, ProfT p) = (m,) . ProfT $! M.alter f i p
   where
-    f Nothing = Just $ singlePath is
-    f (Just q) = Just $ addPath0 is q
+    f Nothing = Just $ singlePath b is
+    f (Just q) = Just $ addPath0 b is q
 
 -- Adds a path to a profile trie, incrementing the count for the given
 -- path.
-addPath :: (Ord k) => [k] -> ProfTrie k Int -> ProfTrie k Int
-addPath [] p = p
-addPath (i : is) (ProfT m) = ProfT $ M.alter f i m
+addPath ::
+  (Ord k) =>
+  Bool ->
+  [k] ->
+  ProfTrie k (Int, Int) ->
+  ProfTrie k (Int, Int)
+addPath _ [] p = p
+addPath b (i : is) (ProfT m) = ProfT $ M.alter f i m
   where
-    f Nothing = Just $ singlePath is
-    f (Just q) = Just $ addPath0 is q
+    f Nothing = Just $ singlePath b is
+    f (Just q) = Just $ addPath0 b is q
 
 data AggInfo k = Ag
   { -- inherited sample count
@@ -116,6 +154,7 @@ prune keep (ProfT m) = case M.traverseMaybeWithKey (prune0 keep) m of
 topN :: (Ord k) => Int -> Map k Int -> [(k, Int)]
 topN n0 = M.foldlWithKey (ins n0) []
   where
+    ins _ pss _ 0 = pss
     ins 0 _ _ _ = []
     ins _ [] k i = [(k, i)]
     ins n pss@((k1, j) : ps) k0 i
@@ -193,9 +232,15 @@ dispProfEntry ppe misc refs (k, ks) (inh, self) =
   where
     ind = fromIntegral $ length ks
 
-dispFunc :: PrettyPrintEnv -> Reference -> Pretty ColorText
-dispFunc ppe =
-  syntaxToColor . prettyHashQualified . termName ppe . Ref
+dispFunc ::
+  PrettyPrintEnv ->
+  Map Reference (Pretty ColorText) ->
+  Reference ->
+  Pretty ColorText
+dispFunc ppe misc r
+  | Just pr <- M.lookup r misc = pr
+  | otherwise =
+      syntaxToColor . prettyHashQualified . termName ppe $ Ref r
 
 dispKey ::
   (Ord k) =>
@@ -205,9 +250,7 @@ dispKey ::
   k ->
   Pretty ColorText
 dispKey ppe misc refs k = case M.lookup k refs of
-  Just r
-    | Just pr <- M.lookup r misc -> pr
-    | otherwise -> dispFunc ppe r
+  Just r -> dispFunc ppe misc r
   Nothing -> "<unknown>"
 
 dispProfTrie ::
@@ -223,10 +266,11 @@ dispProfTrie ppe misc refs ag =
 dispTopEntry ::
   (Ord k) =>
   PrettyPrintEnv ->
+  Map Reference (Pretty ColorText) ->
   Map k Reference ->
   (k, Double) ->
   Pretty ColorText
-dispTopEntry ppe refs (k, frac) =
+dispTopEntry ppe misc refs (k, frac) =
   mconcat
     [ P.indentN 3 . fromString $ showPercent frac,
       P.indentN 4 dr,
@@ -235,16 +279,24 @@ dispTopEntry ppe refs (k, frac) =
   where
     dr :: Pretty ColorText
     dr
-      | Just r <- M.lookup k refs = dispFunc ppe r
+      | Just r <- M.lookup k refs = dispFunc ppe misc r
       | otherwise = "<unknown>"
 
 dispTop ::
   (Ord k) =>
   PrettyPrintEnv ->
+  Map Reference (Pretty ColorText) ->
   Map k Reference ->
   [(k, Double)] ->
   Pretty ColorText
-dispTop ppe refs = foldMap (dispTopEntry ppe refs)
+dispTop ppe misc refs = foldMap (dispTopEntry ppe misc refs)
+
+overallHeader :: Pretty ColorText -> Int -> Pretty ColorText
+overallHeader label samps = label <> ": " <> dsamps <> newline
+  where
+    dsamps
+      | samps == 1 = "1 sample"
+      | otherwise = fromString (show samps) <> " samples"
 
 profileTopHeader :: Pretty ColorText
 profileTopHeader =
@@ -264,38 +316,64 @@ miniProfile ::
   Map Reference (Pretty ColorText) ->
   Profile k ->
   Pretty ColorText
-miniProfile ppe misc (Prof total tr refs) =
-  profileTreeHeader
-    <> dispProfTrie ppe misc refs ag
+miniProfile ppe misc (Prof (total, wtotal) tr refs) =
+  P.lines
+    [ overallHeader "Complete Profile" total,
+      profileTreeHeader <> dispProfTrie ppe misc refs ag,
+      "",
+      if wtotal > 0
+        then
+          overallHeader "Post-wakeup Profile" wtotal
+            <> newline
+            <> profileTreeHeader
+            <> dispProfTrie ppe misc refs agw
+        else "Threads never missed ticks"
+    ]
   where
-    ag = aggregatePruned total tr
+    (full, wait) = demux tr
+    ag = aggregatePruned total full
+    agw = aggregatePruned wtotal wait
 
 fullProfile ::
   (Ord k) =>
   PrettyPrintEnv ->
   Map Reference (Pretty ColorText) ->
   Profile k ->
-  Pretty ColorText
-fullProfile ppe misc (Prof total tr refs) =
-  profileTopHeader
-    <> dispTop ppe refs top
-    <> "\n\n"
-    <> profileTreeHeader
-    <> dispProfTrie ppe misc refs ag
+  (Pretty ColorText, Pretty ColorText)
+fullProfile ppe misc (Prof (total, wtotal) tr0 refs) =
+  ( make "Complete Profile" total comp,
+    make "Post-wakeup Profile" wtotal wait
+  )
   where
-    (top, ag) = aggregate total tr
+    (comp, wait) = demux tr0
+
+    make label tot tr =
+      overallHeader label tot
+        <> newline
+        <> profileTopHeader
+        <> dispTop ppe misc refs top
+        <> newline
+        <> newline
+        <> profileTreeHeader
+        <> dispProfTrie ppe misc refs ag
+      where
+        (top, ag) = aggregate tot tr
 
 foldedProfile ::
   (Ord k) =>
   PrettyPrintEnv ->
   Map Reference (Pretty ColorText) ->
   Profile k ->
-  String
+  (String, String)
 foldedProfile ppe misc (Prof _ tr refs) =
-  toPlain 0 $ foldMapTrie f tr
+  ( toPlain 0 $ foldMapTrie f comp,
+    toPlain 0 $ foldMapTrie f wake
+  )
   where
     dk = dispKey ppe misc refs
 
+    (comp, wake) = demux tr
+
     f (k, ks) n =
       mconcat
         [ foldl (\tx k -> dk k <> ";" <> tx) (dk k) ks,

unison-runtime/src/Unison/Runtime/Interface.hs

@@ -545,10 +545,14 @@ profileEval actThr cleanThr ctxVar cl ppe mout tm = do
         case mout of
           Just loc
             | ticky $ takeExtension loc -> do
-                writeFile loc $ foldedProfile ppe fnames pout
+                let (comp, wake) = foldedProfile ppe fnames pout
+                writeFile loc comp
+                writeFile (loc <.> "wakeup") wake
                 pure $ Right (errs, tmr)
             | otherwise -> do
-                writeFile loc . toPlain 0 $ fullProfile ppe fnames pout
+                let (comp, wake) = fullProfile ppe fnames pout
+                writeFile loc $ toPlain 0 comp
+                writeFile (loc <.> "wakeup") $ toPlain 0 wake
                 pure $ Right (errs, tmr)
           Nothing ->
             pure $ Right (errs <> Profile (miniProfile ppe fnames pout), tmr)

unison-runtime/src/Unison/Runtime/Machine.hs

@@ -568,7 +568,8 @@ eval' !yld env henv !activeThreads !stk !k r (RMatch i pu br) = do
         | otherwise -> unhandledAbilityRequest
 eval' !yld env henv !activeThreads !stk !k here (Yield args)
   | asize stk > 0,
-    VArg1 i <- args =
+    VArg1 i <- args = do
+      checkTicker yld here k
       peekOff stk i >>= apply yld env henv activeThreads stk k False ZArgs
   | otherwise = do
       checkTicker yld here k

unison-runtime/src/Unison/Runtime/Machine/Types.hs

@@ -2,16 +2,16 @@
 
 module Unison.Runtime.Machine.Types where
 
-#if !defined(mingw32_HOST_OS)
-import Control.Concurrent
-  (ThreadId, MVar, newEmptyMVar, tryPutMVar, tryTakeMVar)
-#else
 import Control.Concurrent (ThreadId)
-#endif
-
 import Control.Concurrent.STM as STM
 import Control.Exception hiding (Handler)
-import Data.IORef (IORef, newIORef, readIORef, writeIORef)
+#if !defined(mingw32_HOST_OS)
+import Data.IORef
+  (IORef, newIORef, readIORef, writeIORef, atomicModifyIORef)
+#else
+import Data.IORef
+  (IORef, newIORef, readIORef, writeIORef)
+#endif
 import Data.Kind (Type)
 import Data.Map.Strict qualified as M
 import Data.Set qualified as S
@@ -103,19 +103,20 @@ type Tick = CombIx -> K -> IO ()
 #if !defined(mingw32_HOST_OS)
 -- GHC.Event, time-baed profiler
 instance RuntimeProfiler ProfileComm where
-  newtype Ticker ProfileComm = ProfTicker (MVar Tick)
+  newtype Ticker ProfileComm = ProfTicker (IORef (Maybe Tick))
 
   startTicker (PC pf _ _) = do
-    ticker <- newEmptyMVar
+    ticker <- newIORef Nothing
     cancel <- newIORef False
     tm <- getSystemTimerManager
     void . registerTimeout tm 100 $
       tickCallback 100 pf ticker cancel
     pure (ProfTicker ticker, writeIORef cancel True)
 
-  checkTicker (ProfTicker tick) cix k = tryTakeMVar tick >>= \case
-    Nothing -> pure ()
-    Just pf -> pf cix k
+  checkTicker (ProfTicker ticker) cix k =
+    atomicModifyIORef ticker (Nothing,) >>= \case
+      Nothing -> pure ()
+      Just pf -> pf cix k
   {-# INLINE checkTicker #-}
 
 -- Callback for producing ticks via event manager timeouts. These happen
@@ -126,27 +127,36 @@ instance RuntimeProfiler ProfileComm where
 --
 -- The callback doesn't block trying to write to the MVar, so if something
 -- is already there, a second tick just won't happen.
-tickCallback :: Int -> Tick -> MVar Tick -> IORef Bool -> IO ()
-tickCallback interval tick ticker cancel = body
+tickCallback ::
+  Int ->
+  (Bool -> Tick) ->
+  IORef (Maybe Tick) ->
+  IORef Bool ->
+  IO ()
+tickCallback interval ptick ticker cancel = body
   where
     body = do
-      tryPutMVar ticker tick
+      _full <- atomicModifyIORef ticker \(isJust -> b) ->
+        (Just $ ptick b, b)
       b <- readIORef cancel
       when (not b) do
         tm <- getSystemTimerManager
         () <$ registerTimeout tm interval body
+
 #else
+
 -- CPUTime based profiler for Windows
 instance RuntimeProfiler ProfileComm where
   data Ticker ProfileComm = TPC !Tick !(IORef Word8)
-  startTicker (PC pf _ _) = (, pure ()) . TPC pf  <$> newIORef 1
+  startTicker (PC pf _ _) = (, pure ()) . TPC (pf False)  <$> newIORef 1
 
   checkTicker (TPC tick r) cix k = do
     n <- readIORef r
     when (n `mod` 128 == 0) do
       n <- getCPUTime
       when (n `mod` 100000 == 0) $ tick cix k
     writeIORef r (n+1)
+
 #endif
 
 -- code caching environment