Profiling: Allow heap profiling to be controlled dynamically.

This patch exposes three new functions in `GHC.Profiling` which allow
heap profiling to be enabled and disabled dynamically.

1. startHeapProfTimer - Starts heap profiling with the given RTS options
2. stopHeapProfTimer  - Stops heap profiling
3. requestHeapCensus  - Perform a heap census on the next context
                        switch, regardless of whether the timer is enabled or not.

17 files changed, 182 insertions, 13 deletions

diff --git a/docs/users_guide/9.2.1-notes.rst b/docs/users_guide/9.2.1-notes.rst
index 918f8ebae8..2630ae91fe 100644
--- a/docs/users_guide/9.2.1-notes.rst
+++ b/docs/users_guide/9.2.1-notes.rst
@@ -93,6 +93,12 @@ Compiler
   that the compiler automatically insert cost-centres on all call-sites of
   the named function.
 
+- The heap profiler can now be controlled from within a Haskell program using
+  functions in ``GHC.Profiling``. Profiling can be started and stopped or a heap
+  census requested at a specific point in the program.
+  There is a new RTS flag :rts-flag:`--no-automatic-heap-samples` which can be
+  used to stop heap profiling starting when a program starts.
+
 GHCi
 ~~~~
 
diff --git a/docs/users_guide/profiling.rst b/docs/users_guide/profiling.rst
index 1eddcb7160..6eb5008146 100644
--- a/docs/users_guide/profiling.rst
+++ b/docs/users_guide/profiling.rst
@@ -459,7 +459,7 @@ compiled program.
     :type: dynamic
 
     Deprecated alias for :ghc-flag:`-fprof-auto-exported`
-    
+
 .. ghc-flag:: -caf-all
     :shortdesc: *(deprecated)* Alias for :ghc-flag:`-fprof-cafs`
     :type: dynamic
@@ -886,6 +886,14 @@ There are three more options which relate to heap profiling:
     profiles are always sampled with the frequency of the RTS clock. See
     :ref:`prof-time-options` for changing that.
 
+.. rts-flag:: --no-automatic-heap-samples
+
+    :since: 9.2.1
+
+    Don't start heap profiling from the start of program execution. If this
+    option is enabled, it's expected that the user will manually start heap
+    profiling or request specific samples using functions from ``GHC.Profiling``.
+
 .. rts-flag:: -L ⟨num⟩
 
     Sets the maximum length of a cost-centre stack name in a heap
diff --git a/includes/Rts.h b/includes/Rts.h
index e093a4bcde..50a3f665de 100644
--- a/includes/Rts.h
+++ b/includes/Rts.h
@@ -194,6 +194,7 @@ void _assertFail(const char *filename, unsigned int linenum)
 
 /* Profiling information */
 #include "rts/prof/CCS.h"
+#include "rts/prof/Heap.h"
 #include "rts/prof/LDV.h"
 
 /* Parallel information */
diff --git a/includes/rts/Flags.h b/includes/rts/Flags.h
index 735605b2ba..204ec525ac 100644
--- a/includes/rts/Flags.h
+++ b/includes/rts/Flags.h
@@ -144,6 +144,7 @@ typedef struct _PROFILING_FLAGS {
 
     Time        heapProfileInterval; /* time between samples */
     uint32_t    heapProfileIntervalTicks; /* ticks between samples (derived) */
+    bool        startHeapProfileAtStartup; /* true if we start profiling from program startup */
 
 
     bool		showCCSOnException;
diff --git a/includes/rts/prof/Heap.h b/includes/rts/prof/Heap.h
new file mode 100644
index 0000000000..90700c809b
--- /dev/null
+++ b/includes/rts/prof/Heap.h
@@ -0,0 +1,24 @@
+/* -----------------------------------------------------------------------------
+ *
+ * (c) The University of Glasgow, 2009
+ *
+ * Heap Census Profiling
+ *
+ * Do not #include this file directly: #include "Rts.h" instead.
+ *
+ * To understand the structure of the RTS headers, see the wiki:
+ *   https://gitlab.haskell.org/ghc/ghc/wikis/commentary/source-tree/includes
+ *
+ * ---------------------------------------------------------------------------*/
+
+#pragma once
+
+/* -----------------------------------------------------------------------------
+ * Fine-grained control over heap census profiling which can be called from
+ * Haskell to restrict the profile to portion(s) of the execution.
+ * See the module GHC.Profiling.
+ * ---------------------------------------------------------------------------*/
+
+void requestHeapCensus ( void );
+void startHeapProfTimer ( void );
+void stopHeapProfTimer ( void );
diff --git a/libraries/base/GHC/Profiling.hs b/libraries/base/GHC/Profiling.hs
index 917a208b30..b7cef2fecd 100644
--- a/libraries/base/GHC/Profiling.hs
+++ b/libraries/base/GHC/Profiling.hs
@@ -2,7 +2,14 @@
 {-# LANGUAGE NoImplicitPrelude #-}
 
 -- | @since 4.7.0.0
-module GHC.Profiling where
+module GHC.Profiling ( -- * Cost Centre Profiling
+                       startProfTimer
+                     , stopProfTimer
+                       -- * Heap Profiling
+                     , startHeapProfTimer
+                     , stopHeapProfTimer
+                     , requestHeapCensus
+                     )where
 
 import GHC.Base
 
@@ -17,3 +24,26 @@ foreign import ccall stopProfTimer :: IO ()
 --
 -- @since 4.7.0.0
 foreign import ccall startProfTimer :: IO ()
+
+-- | Request a heap census on the next context switch. The census can be
+-- requested whether or not the heap profiling timer is running.
+-- Note: This won't do anything unless you also specify a profiling mode on the
+-- command line using the normal RTS options.
+--
+-- @since 4.16.0.0
+foreign import ccall requestHeapCensus :: IO ()
+
+-- | Start heap profiling. This is called normally by the RTS on start-up,
+-- but can be disabled using the rts flag `--no-automatic-heap-samples`
+-- Note: This won't do anything unless you also specify a profiling mode on the
+-- command line using the normal RTS options.
+--
+-- @since 4.16.0.0
+foreign import ccall startHeapProfTimer :: IO ()
+
+-- | Stop heap profiling.
+-- Note: This won't do anything unless you also specify a profiling mode on the
+-- command line using the normal RTS options.
+--
+-- @since 4.16.0.0
+foreign import ccall stopHeapProfTimer :: IO ()
diff --git a/libraries/base/GHC/RTS/Flags.hsc b/libraries/base/GHC/RTS/Flags.hsc
index 2130ee1aaa..2f161c72c3 100644
--- a/libraries/base/GHC/RTS/Flags.hsc
+++ b/libraries/base/GHC/RTS/Flags.hsc
@@ -288,6 +288,7 @@ data ProfFlags = ProfFlags
     { doHeapProfile            :: DoHeapProfile
     , heapProfileInterval      :: RtsTime -- ^ time between samples
     , heapProfileIntervalTicks :: Word    -- ^ ticks between samples (derived)
+    , startHeapProfileAtStartup :: Bool
     , showCCSOnException       :: Bool
     , maxRetainerSetSize       :: Word
     , ccsLength                :: Word
@@ -584,6 +585,8 @@ getProfFlags = do
             <*> #{peek PROFILING_FLAGS, heapProfileInterval} ptr
             <*> #{peek PROFILING_FLAGS, heapProfileIntervalTicks} ptr
             <*> (toBool <$>
+                  (#{peek PROFILING_FLAGS, startHeapProfileAtStartup} ptr :: IO CBool))
+            <*> (toBool <$>
                   (#{peek PROFILING_FLAGS, showCCSOnException} ptr :: IO CBool))
             <*> #{peek PROFILING_FLAGS, maxRetainerSetSize} ptr
             <*> #{peek PROFILING_FLAGS, ccsLength} ptr
diff --git a/rts/Proftimer.c b/rts/Proftimer.c
index 00b92a227d..29abb62a8e 100644
--- a/rts/Proftimer.c
+++ b/rts/Proftimer.c
@@ -18,7 +18,12 @@
 static bool do_prof_ticks = false;       // enable profiling ticks
 #endif
 
-static bool do_heap_prof_ticks = false;  // enable heap profiling ticks
+static bool do_heap_prof_ticks = false;  // Whether the timer is currently ticking down
+static bool heap_prof_timer_active = false;  // Whether the timer is enabled at all
+
+/* The heap_prof_timer_active flag controls whether heap profiling is enabled
+at all, once it is enabled, the `do_heap_prof_ticks` flag controls whether the
+counter is currently counting down. This is paused, for example, in Schedule.c. */
 
 // Sampling of Ticky-Ticky profiler to eventlog
 #if defined(TICKY_TICKY) && defined(TRACING)
@@ -51,26 +56,56 @@ startProfTimer( void )
 void
 stopHeapProfTimer( void )
 {
-    RELAXED_STORE(&do_heap_prof_ticks, false);
+  if (RtsFlags.ProfFlags.doHeapProfile){
+    RELAXED_STORE(&heap_prof_timer_active, false);
+    pauseHeapProfTimer();
+  }
 }
 
 void
 startHeapProfTimer( void )
 {
+  if (RtsFlags.ProfFlags.doHeapProfile){
+    RELAXED_STORE(&heap_prof_timer_active, true);
+    resumeHeapProfTimer();
+  }
+}
+
+void
+pauseHeapProfTimer ( void ) {
+    RELAXED_STORE(&do_heap_prof_ticks, false);
+}
+
+
+void
+resumeHeapProfTimer ( void ) {
     if (RtsFlags.ProfFlags.doHeapProfile &&
         RtsFlags.ProfFlags.heapProfileIntervalTicks > 0) {
-        do_heap_prof_ticks = true;
+        RELAXED_STORE(&do_heap_prof_ticks, true);
     }
 }
 
 void
+requestHeapCensus( void ){
+  // If no profiling mode is passed then just ignore the call.
+  if (RtsFlags.ProfFlags.doHeapProfile){
+    RELAXED_STORE(&performHeapProfile, true);
+  }
+}
+
+void
 initProfTimer( void )
 {
     performHeapProfile = false;
 
     ticks_to_heap_profile = RtsFlags.ProfFlags.heapProfileIntervalTicks;
 
-    startHeapProfTimer();
+    /* This might look a bit strange but the heap profile timer can
+      be toggled on/off from within Haskell by calling the startHeapProf
+      function from within Haskell */
+    if (RtsFlags.ProfFlags.startHeapProfileAtStartup){
+      startHeapProfTimer();
+    }
 }
 
 uint32_t total_ticks = 0;
@@ -99,7 +134,7 @@ handleProfTick(void)
     }
 #endif
 
-    if (RELAXED_LOAD(&do_heap_prof_ticks)) {
+    if (RELAXED_LOAD(&do_heap_prof_ticks) && RELAXED_LOAD(&heap_prof_timer_active))  {
         ticks_to_heap_profile--;
         if (ticks_to_heap_profile <= 0) {
             ticks_to_heap_profile = RtsFlags.ProfFlags.heapProfileIntervalTicks;
diff --git a/rts/Proftimer.h b/rts/Proftimer.h
index ad50ccb9a6..f4432e46bc 100644
--- a/rts/Proftimer.h
+++ b/rts/Proftimer.h
@@ -12,9 +12,8 @@
 
 void initProfTimer      ( void );
 void handleProfTick     ( void );
-
-void stopHeapProfTimer  ( void );
-void startHeapProfTimer ( void );
+void pauseHeapProfTimer  ( void );
+void resumeHeapProfTimer ( void );
 
 extern bool performHeapProfile;
 extern bool performTickySample;
diff --git a/rts/RtsFlags.c b/rts/RtsFlags.c
index a5ec423f8d..bc7e86901f 100644
--- a/rts/RtsFlags.c
+++ b/rts/RtsFlags.c
@@ -212,6 +212,7 @@ void initRtsFlagsDefaults(void)
 
     RtsFlags.ProfFlags.doHeapProfile      = false;
     RtsFlags.ProfFlags.heapProfileInterval = USToTime(100000); // 100ms
+    RtsFlags.ProfFlags.startHeapProfileAtStartup = true;
 
 #if defined(PROFILING)
     RtsFlags.ProfFlags.showCCSOnException = false;
@@ -391,6 +392,11 @@ usage_text[] = {
 "  -hT      Produce a heap profile grouped by closure type",
 #endif /* PROFILING */
 
+"  -i<sec>  Time between heap profile samples (seconds, default: 0.1)",
+"  --no-automatic-heap-samples  Do not start the heap profile interval timer on start-up,",
+"                               Rather, the application will be responsible for triggering",
+"                               heap profiler samples."
+
 #if defined(TRACING)
 "",
 "  -ol<file>  Send binary eventlog to <file> (default: <program>.eventlog)",
@@ -416,7 +422,6 @@ usage_text[] = {
 "             the initial enabled event classes are 'sgpu'",
 #endif
 
-"  -i<sec>  Time between heap profile samples (seconds, default: 0.1)",
 "",
 #if defined(TICKY_TICKY)
 "  -r<file>  Produce ticky-ticky statistics (with -rstderr for stderr)",
@@ -1091,6 +1096,12 @@ error = true;
                       }
                       break;
                   }
+                  else if (strequal("no-automatic-heap-samples",
+                               &rts_argv[arg][2])) {
+                      OPTION_SAFE;
+                      RtsFlags.ProfFlags.startHeapProfileAtStartup = false;
+                      break;
+                  }
                   else {
                       OPTION_SAFE;
                       errorBelch("unknown RTS option: %s",rts_argv[arg]);
diff --git a/rts/RtsSymbols.c b/rts/RtsSymbols.c
index 7ea833ce55..989b878b56 100644
--- a/rts/RtsSymbols.c
+++ b/rts/RtsSymbols.c
@@ -984,6 +984,9 @@
       SymI_HasProto(unlockFile)                                         \
       SymI_HasProto(startProfTimer)                                     \
       SymI_HasProto(stopProfTimer)                                      \
+      SymI_HasProto(startHeapProfTimer)                                 \
+      SymI_HasProto(stopHeapProfTimer)                                  \
+      SymI_HasProto(requestHeapCensus)                                  \
       SymI_HasProto(atomic_inc)                                         \
       SymI_HasProto(atomic_dec)                                         \
       SymI_HasProto(hs_spt_lookup)                                      \
diff --git a/rts/Schedule.c b/rts/Schedule.c
index 390e505cf2..d9d5c9a74a 100644
--- a/rts/Schedule.c
+++ b/rts/Schedule.c
@@ -416,7 +416,7 @@ run_thread:
     // that.
     cap->r.rCurrentTSO = t;
 
-    startHeapProfTimer();
+    resumeHeapProfTimer();
 
     // ----------------------------------------------------------------------
     // Run the current thread
@@ -534,7 +534,7 @@ run_thread:
     // ----------------------------------------------------------------------
 
     // Costs for the scheduler are assigned to CCS_SYSTEM
-    stopHeapProfTimer();
+    pauseHeapProfTimer();
 #if defined(PROFILING)
     cap->r.rCCCS = CCS_SYSTEM;
 #endif
diff --git a/rts/rts.cabal.in b/rts/rts.cabal.in
index a1d0ce39a2..108b4d6b9b 100644
--- a/rts/rts.cabal.in
+++ b/rts/rts.cabal.in
@@ -180,6 +180,7 @@ library
                       rts/Types.h
                       rts/Utils.h
                       rts/prof/CCS.h
+                      rts/prof/Heap.h
                       rts/prof/LDV.h
                       rts/storage/Block.h
                       rts/storage/ClosureMacros.h
diff --git a/testsuite/tests/profiling/should_run/all.T b/testsuite/tests/profiling/should_run/all.T
index fbe1379e92..9f1fa67e1e 100644
--- a/testsuite/tests/profiling/should_run/all.T
+++ b/testsuite/tests/profiling/should_run/all.T
@@ -8,6 +8,12 @@ test('heapprof002',
 
 test('T11489', [req_profiling], makefile_test, ['T11489'])
 
+test('dynamic-prof', [], compile_and_run, [''])
+
+test('dynamic-prof2', [only_ways(['normal']), extra_run_opts('+RTS -hT --no-automatic-heap-samples')], compile_and_run, [''])
+
+test('dynamic-prof3', [only_ways(['normal']), extra_run_opts('+RTS -hT --no-automatic-heap-samples')], compile_and_run, [''])
+
 # Below this line, run tests only with profiling ways.
 setTestOpts(req_profiling)
 setTestOpts(extra_ways(['prof', 'ghci-ext-prof']))
diff --git a/testsuite/tests/profiling/should_run/dynamic-prof.hs b/testsuite/tests/profiling/should_run/dynamic-prof.hs
new file mode 100644
index 0000000000..243e094877
--- /dev/null
+++ b/testsuite/tests/profiling/should_run/dynamic-prof.hs
@@ -0,0 +1,13 @@
+{-# LANGUAGE BangPatterns #-}
+module Main where
+
+import GHC.Profiling
+import Control.Exception
+
+main = do
+  let !t = [0..1000000]
+  evaluate (length t)
+  requestHeapCensus
+  evaluate (length t)
+
+
diff --git a/testsuite/tests/profiling/should_run/dynamic-prof2.hs b/testsuite/tests/profiling/should_run/dynamic-prof2.hs
new file mode 100644
index 0000000000..243e094877
--- /dev/null
+++ b/testsuite/tests/profiling/should_run/dynamic-prof2.hs
@@ -0,0 +1,13 @@
+{-# LANGUAGE BangPatterns #-}
+module Main where
+
+import GHC.Profiling
+import Control.Exception
+
+main = do
+  let !t = [0..1000000]
+  evaluate (length t)
+  requestHeapCensus
+  evaluate (length t)
+
+
diff --git a/testsuite/tests/profiling/should_run/dynamic-prof3.hs b/testsuite/tests/profiling/should_run/dynamic-prof3.hs
new file mode 100644
index 0000000000..f5a17bef8d
--- /dev/null
+++ b/testsuite/tests/profiling/should_run/dynamic-prof3.hs
@@ -0,0 +1,15 @@
+{-# LANGUAGE BangPatterns #-}
+module Main where
+
+import GHC.Profiling
+import Control.Exception
+
+main = do
+  startHeapProfTimer
+  let !t = [0..1000000]
+  evaluate (length t)
+  requestHeapCensus
+  evaluate (length t)
+  stopHeapProfTimer
+
+