Create option for unison that will write RTS stats to a file

fixes #3597

unison-cli/package.yaml

@@ -118,7 +118,7 @@ executables:
         other-modules: Paths_unison_cli
     source-dirs: unison
     main: Main.hs
-    ghc-options: -threaded -rtsopts "-with-rtsopts=-I0 -N -qn4 -qg1" -optP-Wno-nonportable-include-path
+    ghc-options: -threaded -rtsopts "-with-rtsopts=-I0 -N -qn4 -qg1 -T" -optP-Wno-nonportable-include-path
     dependencies:
       - code-page
       - optparse-applicative >= 0.16.1.0

unison-cli/unison-cli.cabal

@@ -451,6 +451,7 @@ executable unison
   main-is: Main.hs
   other-modules:
       ArgParse
+      Stats
       System.Path
       Version
   hs-source-dirs:
@@ -485,7 +486,7 @@ executable unison
       TupleSections
       TypeApplications
       ViewPatterns
-  ghc-options: -Wall -threaded -rtsopts "-with-rtsopts=-I0 -N -qn4 -qg1" -optP-Wno-nonportable-include-path
+  ghc-options: -Wall -threaded -rtsopts "-with-rtsopts=-I0 -N -qn4 -qg1 -T" -optP-Wno-nonportable-include-path
   build-depends:
       IntervalMap
     , ListLike

unison-cli/unison/ArgParse.hs

@@ -53,6 +53,7 @@ import qualified Options.Applicative as OptParse
 import Options.Applicative.Builder.Internal (noGlobal {- https://github.com/pcapriotti/optparse-applicative/issues/461 -})
 import Options.Applicative.Help (bold, (<+>))
 import qualified Options.Applicative.Help.Pretty as P
+import Stats
 import System.Environment (lookupEnv)
 import Text.Read (readMaybe)
 import qualified Unison.Codebase.Path as Path
@@ -116,7 +117,7 @@ data Command
   | -- @deprecated in trunk after M2g. Remove the Init command completely after M2h has been released
     Init
   | Run RunSource [String]
-  | Transcript ShouldForkCodebase ShouldSaveCodebase (NonEmpty FilePath)
+  | Transcript ShouldForkCodebase ShouldSaveCodebase (Maybe RtsStatsPath) (NonEmpty FilePath)
   deriving (Show, Eq)
 
 -- | Options shared by sufficiently many subcommands.
@@ -381,6 +382,15 @@ runCompiledParser :: Parser Command
 runCompiledParser =
   Run . RunCompiled <$> fileArgument "path/to/file" <*> runArgumentParser
 
+rtsStatsOption :: Parser (Maybe RtsStatsPath)
+rtsStatsOption =
+  let meta =
+        metavar "FILE.json"
+          <> long "rts-stats"
+          <> help "Write json summary of rts stats to FILE"
+          <> noGlobal
+   in optional (option OptParse.str meta)
+
 saveCodebaseFlag :: Parser ShouldSaveCodebase
 saveCodebaseFlag = flag DontSaveCodebase SaveCodebase (long "save-codebase" <> help saveHelp)
   where
@@ -448,15 +458,17 @@ transcriptParser :: Parser Command
 transcriptParser = do
   -- ApplicativeDo
   shouldSaveCodebase <- saveCodebaseFlag
+  mrtsStatsFp <- rtsStatsOption
   files <- liftA2 (NE.:|) (fileArgument "FILE") (many (fileArgument "FILES..."))
-  pure (Transcript DontFork shouldSaveCodebase files)
+  pure (Transcript DontFork shouldSaveCodebase mrtsStatsFp files)
 
 transcriptForkParser :: Parser Command
 transcriptForkParser = do
   -- ApplicativeDo
   shouldSaveCodebase <- saveCodebaseFlag
+  mrtsStatsFp <- rtsStatsOption
   files <- liftA2 (NE.:|) (fileArgument "FILE") (many (fileArgument "FILES..."))
-  pure (Transcript UseFork shouldSaveCodebase files)
+  pure (Transcript UseFork shouldSaveCodebase mrtsStatsFp files)
 
 unisonHelp :: String -> String -> P.Doc
 unisonHelp (P.text -> executable) (P.text -> version) =

unison-cli/unison/Main.hs

@@ -42,6 +42,7 @@ import qualified Language.Haskell.TH as TH
 import qualified Language.Haskell.TH.Syntax as TH
 import qualified Network.HTTP.Client as HTTP
 import qualified Network.HTTP.Client.TLS as HTTP
+import Stats (recordRtsStats)
 import System.Directory (canonicalizePath, getCurrentDirectory, removeDirectoryRecursive)
 import System.Environment (getProgName, lookupEnv, withArgs)
 import qualified System.Exit as Exit
@@ -246,8 +247,11 @@ main = withCP65001 . runInUnboundThread . Ki.scoped $ \scope -> do
                         "to produce a new compiled program \
                         \that matches your version of Unison."
                     ]
-      Transcript shouldFork shouldSaveCodebase transcriptFiles ->
-        runTranscripts renderUsageInfo shouldFork shouldSaveCodebase mCodePathOption transcriptFiles
+      Transcript shouldFork shouldSaveCodebase mrtsStatsFp transcriptFiles -> do
+        let action = runTranscripts renderUsageInfo shouldFork shouldSaveCodebase mCodePathOption transcriptFiles
+        case mrtsStatsFp of
+          Nothing -> action
+          Just fp -> recordRtsStats fp action
       Launch isHeadless codebaseServerOpts downloadBase mayStartingPath shouldWatchFiles -> do
         getCodebaseOrExit mCodePathOption SC.MigrateAfterPrompt \(initRes, _, theCodebase) -> do
           runtime <- RTI.startRuntime False RTI.Persistent Version.gitDescribeWithDate

unison-cli/unison/Stats.hs

@@ -0,0 +1,39 @@
+module Stats
+  ( RtsStatsPath (..),
+    recordRtsStats,
+  )
+where
+
+import Control.Exception (finally)
+import Data.Aeson (encode, object, (.=))
+import qualified Data.ByteString.Lazy as BL
+import Data.Function
+import Data.String (IsString)
+import GHC.Stats
+
+newtype RtsStatsPath
+  = RtsStatsPath FilePath
+  deriving stock (Show, Eq)
+  deriving newtype (IsString)
+
+recordRtsStats :: RtsStatsPath -> IO a -> IO a
+recordRtsStats (RtsStatsPath fp) action = do
+  r0 <- getRTSStats
+  action `finally` do
+    r1 <- getRTSStats
+    BL.writeFile fp (encode (produceStats r0 r1))
+  where
+    produceStats r0 r1 =
+      object
+        [ "gcs" .= on (-) gcs r1 r0,
+          "major_gcs" .= on (-) major_gcs r1 r0,
+          "allocated_bytes" .= on (-) allocated_bytes r1 r0,
+          "max_live_bytes" .= on (-) max_live_bytes r1 r0,
+          "copied_bytes" .= on (-) copied_bytes r1 r0,
+          "mutator_cpu_ns" .= on (-) mutator_cpu_ns r1 r0,
+          "mutator_elapsed_ns" .= on (-) mutator_elapsed_ns r1 r0,
+          "gc_cpu_ns" .= on (-) mutator_cpu_ns r1 r0,
+          "gc_elapsed_ns" .= on (-) mutator_elapsed_ns r1 r0,
+          "cpu_ns" .= on (-) cpu_ns r1 r0,
+          "elapsed_ns" .= on (-) cpu_ns r1 r0
+        ]