Emit final heap alloc events and rearrange code to calculate alloc totals

In stat_exit we want to emit a final EVENT_HEAP_ALLOCATED for each cap
so that we get the same total allocation count as reported via +RTS -s.
To do so we need to update the per-cap total_allocated counts.

Previously we had a single calcAllocated(rtsBool) function that counted
the large allocations and optionally the nurseries for all caps. The GC
would always call it with false, and the stat_exit always with true.
The reason for these two modes is that the GC counts the nurseries via
clearNurseries() (which also updates the per-cap total_allocated
counts), so it's only the stat_exit() path that needs to count them.

We now split the calcAllocated() function into two: countLargeAllocated
and updateNurseriesStats. As the name suggests, the latter now updates
the per-cap total_allocated counts, in additon to returning a total.

4 files changed, 34 insertions, 24 deletions

diff --git a/rts/Stats.c b/rts/Stats.c
index a7335531ce..29cae21bf3 100644
--- a/rts/Stats.c
+++ b/rts/Stats.c
@@ -588,6 +588,12 @@ stat_exit(int alloc)
 
 	GC_tot_alloc += alloc;
 
+        for (i = 0; i < n_capabilities; i++) {
+            traceEventHeapAllocated(&capabilities[i],
+                                    CAPSET_HEAP_DEFAULT,
+                                    capabilities[i].total_allocated * sizeof(W_));
+        }
+
 	/* Count total garbage collections */
 	for (g = 0; g < RtsFlags.GcFlags.generations; g++)
 	    total_collections += generations[g].collections;
diff --git a/rts/sm/GC.c b/rts/sm/GC.c
index 055a1363ed..0d83f2abae 100644
--- a/rts/sm/GC.c
+++ b/rts/sm/GC.c
@@ -231,7 +231,7 @@ GarbageCollect (rtsBool force_major_gc,
   /* Approximate how much we allocated.  
    * Todo: only when generating stats? 
    */
-  allocated = calcAllocated(rtsFalse/* don't count the nursery yet */);
+  allocated = countLargeAllocated(); /* don't count the nursery yet */
 
   /* Figure out which generation to collect
    */
diff --git a/rts/sm/Storage.c b/rts/sm/Storage.c
index 4cd1bc99c9..7da0c70e46 100644
--- a/rts/sm/Storage.c
+++ b/rts/sm/Storage.c
@@ -235,7 +235,8 @@ void storageAddCapabilities (nat from, nat to)
 void
 exitStorage (void)
 {
-    stat_exit(calcAllocated(rtsTrue));
+    lnat allocated = updateNurseriesStats();
+    stat_exit(allocated);
 }
 
 void
@@ -904,34 +905,36 @@ dirty_MVAR(StgRegTable *reg, StgClosure *p)
  * -------------------------------------------------------------------------- */
 
 /* -----------------------------------------------------------------------------
- * calcAllocated()
+ * updateNurseriesStats()
  *
- * Approximate how much we've allocated: number of blocks in the
- * nursery + blocks allocated via allocate() - unused nusery blocks.
- * This leaves a little slop at the end of each block.
+ * Update the per-cap total_allocated numbers with an approximation of
+ * the amount of memory used in each cap's nursery. Also return the
+ * total across all caps.
+ * 
+ * Since this update is also performed by clearNurseries() then we only
+ * need this function for the final stats when the RTS is shutting down.
  * -------------------------------------------------------------------------- */
 
 lnat
-calcAllocated (rtsBool include_nurseries)
+updateNurseriesStats (void)
 {
-  nat allocated = 0;
-  nat i;
-
-  // When called from GC.c, we already have the allocation count for
-  // the nursery from resetNurseries(), so we don't need to walk
-  // through these block lists again.
-  if (include_nurseries)
-  {
-      for (i = 0; i < n_capabilities; i++) {
-          allocated += countOccupied(nurseries[i].blocks);
-      }
-  }
+    lnat allocated = 0;
+    nat i;
 
-  // add in sizes of new large and pinned objects
-  allocated += g0->n_new_large_words;
+    for (i = 0; i < n_capabilities; i++) {
+        int cap_allocated = countOccupied(nurseries[i].blocks);
+        capabilities[i].total_allocated += cap_allocated;
+        allocated                       += cap_allocated;
+    }
 
-  return allocated;
-}  
+    return allocated;
+}
+
+lnat
+countLargeAllocated (void)
+{
+    return g0->n_new_large_words;
+}
 
 lnat countOccupied (bdescr *bd)
 {
diff --git a/rts/sm/Storage.h b/rts/sm/Storage.h
index 11afc26598..44f39ee29b 100644
--- a/rts/sm/Storage.h
+++ b/rts/sm/Storage.h
@@ -90,7 +90,8 @@ lnat     countNurseryBlocks   ( void );
    Stats 'n' DEBUG stuff
    -------------------------------------------------------------------------- */
 
-lnat    calcAllocated  (rtsBool count_nurseries);
+lnat    updateNurseriesStats (void);
+lnat    countLargeAllocated  (void);
 lnat    countOccupied  (bdescr *bd);
 lnat    calcNeeded     (void);