From 3703bbd631d6dff42b7046d1a520dc7f13500933 Mon Sep 17 00:00:00 2001
From: Kyle Katarn <contact@kcsoftwares.com>
Date: Mon, 12 Sep 2022 14:22:52 +0200
Subject: [PATCH] [NTOS:MM] Implement PeakCommitment (MmPeakCommitment,
 MmTotalCommittedPages) (#4650)

And return the corresponding values in SystemPerformanceInformation.
Lockless updating counters suggestion by Thomas Faber.
---
 ntoskrnl/ex/sysinfo.c          | 11 +++--------
 ntoskrnl/include/internal/mm.h | 27 +++++++++++++++++++++++++++
 ntoskrnl/mm/balance.c          |  2 ++
 ntoskrnl/mm/pagefile.c         |  3 +++
 4 files changed, 35 insertions(+), 8 deletions(-)

diff --git a/ntoskrnl/ex/sysinfo.c b/ntoskrnl/ex/sysinfo.c
index bf9b957fd42..ecf6380dd8b 100644
--- a/ntoskrnl/ex/sysinfo.c
+++ b/ntoskrnl/ex/sysinfo.c
@@ -714,13 +714,8 @@ QSI_DEF(SystemPerformanceInformation)
     }
 
     Spi->AvailablePages = (ULONG)MmAvailablePages;
-    /*
-     *   Add up all the used "Committed" memory + pagefile.
-     *   Not sure this is right. 8^\
-     */
-    Spi->CommittedPages = MiMemoryConsumers[MC_SYSTEM].PagesUsed +
-                          MiMemoryConsumers[MC_USER].PagesUsed +
-                          MiUsedSwapPages;
+
+    Spi->CommittedPages = MmTotalCommittedPages;
     /*
      *  Add up the full system total + pagefile.
      *  All this make Taskmgr happy but not sure it is the right numbers.
@@ -728,7 +723,7 @@ QSI_DEF(SystemPerformanceInformation)
      */
     Spi->CommitLimit = MmNumberOfPhysicalPages + MiFreeSwapPages + MiUsedSwapPages;
 
-    Spi->PeakCommitment = 0; /* FIXME */
+    Spi->PeakCommitment = MmPeakCommitment;
     Spi->PageFaultCount = 0; /* FIXME */
     Spi->CopyOnWriteCount = 0; /* FIXME */
     Spi->TransitionCount = 0; /* FIXME */
diff --git a/ntoskrnl/include/internal/mm.h b/ntoskrnl/include/internal/mm.h
index 5c28d27806b..af281b26ff3 100644
--- a/ntoskrnl/include/internal/mm.h
+++ b/ntoskrnl/include/internal/mm.h
@@ -866,6 +866,33 @@ NTAPI
 MmDeleteKernelStack(PVOID Stack,
                     BOOLEAN GuiStack);
 
+/* balance.c / pagefile.c******************************************************/
+
+inline VOID UpdateTotalCommittedPages(LONG Delta)
+{
+    /*
+     * Add up all the used "Committed" memory + pagefile.
+     * Not sure this is right. 8^\
+     * MmTotalCommittedPages should be adjusted consistently with
+     * other counters at different places.
+     *
+       MmTotalCommittedPages = MiMemoryConsumers[MC_SYSTEM].PagesUsed +
+                               MiMemoryConsumers[MC_USER].PagesUsed +
+                               MiUsedSwapPages;
+     */
+    
+    /* Update Commitment */
+    SIZE_T TotalCommittedPages = InterlockedExchangeAddSizeT(&MmTotalCommittedPages, Delta) + Delta;
+
+    /* Update Peak = max(Peak, Total) in a lockless way */
+    SIZE_T PeakCommitment = MmPeakCommitment;
+    while (TotalCommittedPages > PeakCommitment &&
+           InterlockedCompareExchangeSizeT(&MmPeakCommitment, TotalCommittedPages, PeakCommitment) != PeakCommitment)
+    {
+        PeakCommitment = MmPeakCommitment;
+    }
+}
+
 /* balance.c *****************************************************************/
 
 CODE_SEG("INIT")
diff --git a/ntoskrnl/mm/balance.c b/ntoskrnl/mm/balance.c
index 4558deeff18..eb4c9fc8a05 100644
--- a/ntoskrnl/mm/balance.c
+++ b/ntoskrnl/mm/balance.c
@@ -81,6 +81,7 @@ MmReleasePageMemoryConsumer(ULONG Consumer, PFN_NUMBER Page)
     }
 
     (void)InterlockedDecrementUL(&MiMemoryConsumers[Consumer].PagesUsed);
+    UpdateTotalCommittedPages(-1);
 
     OldIrql = MiAcquirePfnLock();
 
@@ -284,6 +285,7 @@ MmRequestPageMemoryConsumer(ULONG Consumer, BOOLEAN CanWait,
 
     /* Update the target */
     InterlockedIncrementUL(&MiMemoryConsumers[Consumer].PagesUsed);
+    UpdateTotalCommittedPages(1);
 
     /*
      * Actually allocate the page.
diff --git a/ntoskrnl/mm/pagefile.c b/ntoskrnl/mm/pagefile.c
index e0f9e087150..6f73c28d37d 100644
--- a/ntoskrnl/mm/pagefile.c
+++ b/ntoskrnl/mm/pagefile.c
@@ -295,6 +295,7 @@ MmFreeSwapPage(SWAPENTRY Entry)
 
     MiFreeSwapPages++;
     MiUsedSwapPages--;
+    UpdateTotalCommittedPages(-1);
 
     KeReleaseGuardedMutex(&MmPageFileCreationLock);
 }
@@ -329,6 +330,8 @@ MmAllocSwapPage(VOID)
             }
             MiUsedSwapPages++;
             MiFreeSwapPages--;
+            UpdateTotalCommittedPages(1);
+
             KeReleaseGuardedMutex(&MmPageFileCreationLock);
 
             entry = ENTRY_FROM_FILE_OFFSET(i, off + 1);