[esan|wset] Iterate all memory to compute the total working set

Summary:
Adds iteration of all application memory in an efficient manner using
shadow faults.  Shadow memory starts out inaccessible and we mark it
writable one page at a time on each fault when the instrumentation touches
it.  This allows iteration over just the mapped shadow memory, saving
significant time.

Adds a process-end iteration and pretty-printing of the final result.

Adds a new test and updates the existing tests.

Reviewers: aizatsky, filcab

Subscribers: vitalybuka, zhaoqin, kcc, eugenis, llvm-commits, kubabrecka

Differential Revision: http://reviews.llvm.org/D20578

git-svn-id: https://llvm.org/svn/llvm-project/compiler-rt/trunk@271277 91177308-0d34-0410-b5e6-96231b3b80d8

3 files changed, 93 insertions, 11 deletions

diff --git a/lib/esan/esan.cpp b/lib/esan/esan.cpp
index c9629d65c..f0a496518 100644
--- a/lib/esan/esan.cpp
+++ b/lib/esan/esan.cpp
@@ -149,8 +149,16 @@ static void initializeShadow() {
     VPrintf(1, "Shadow #%d: [%zx-%zx) (%zuGB)\n", i, ShadowStart, ShadowEnd,
             (ShadowEnd - ShadowStart) >> 30);
 
-    uptr Map = (uptr)MmapFixedNoReserve(ShadowStart, ShadowEnd - ShadowStart,
-                                        "shadow");
+    uptr Map;
+    if (WhichTool == ESAN_WorkingSet) {
+      // We want to identify all shadow pages that are touched so we start
+      // out inaccessible.
+      Map = (uptr)MmapFixedNoAccess(ShadowStart, ShadowEnd- ShadowStart,
+                                    "shadow");
+    } else {
+      Map = (uptr)MmapFixedNoReserve(ShadowStart, ShadowEnd - ShadowStart,
+                                     "shadow");
+    }
     if (Map != ShadowStart) {
       Printf("FATAL: EfficiencySanitizer failed to map its shadow memory.\n");
       Die();
diff --git a/lib/esan/working_set.cpp b/lib/esan/working_set.cpp
index 68ea9e74a..dcd683737 100644
--- a/lib/esan/working_set.cpp
+++ b/lib/esan/working_set.cpp
@@ -16,6 +16,7 @@
 #include "esan.h"
 #include "esan_flags.h"
 #include "esan_shadow.h"
+#include "sanitizer_common/sanitizer_procmaps.h"
 
 // We shadow every cache line of app memory with one shadow byte.
 // - The highest bit of each shadow byte indicates whether the corresponding
@@ -30,6 +31,9 @@ typedef unsigned char byte;
 
 namespace __esan {
 
+// Our shadow memory assumes that the line size is 64.
+static const u32 CacheLineSize = 64;
+
 // See the shadow byte layout description above.
 static const u32 TotalWorkingSetBitIdx = 7;
 static const u32 CurWorkingSetBitIdx = 0;
@@ -75,16 +79,80 @@ void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
   }
 }
 
+// This routine will word-align ShadowStart and ShadowEnd prior to scanning.
+static u32 countAndClearShadowValues(u32 BitIdx, uptr ShadowStart,
+                                     uptr ShadowEnd) {
+  u32 WorkingSetSize = 0;
+  u32 ByteValue = 0x1 << BitIdx;
+  u32 WordValue = ByteValue | ByteValue << 8 | ByteValue << 16 |
+    ByteValue << 24;
+  // Get word aligned start.
+  ShadowStart = RoundDownTo(ShadowStart, sizeof(u32));
+  for (u32 *Ptr = (u32 *)ShadowStart; Ptr < (u32 *)ShadowEnd; ++Ptr) {
+    if ((*Ptr & WordValue) != 0) {
+      byte *BytePtr = (byte *)Ptr;
+      for (u32 j = 0; j < sizeof(u32); ++j) {
+        if (BytePtr[j] & ByteValue) {
+          ++WorkingSetSize;
+          // TODO: Accumulate to the lower-frequency bit to the left.
+        }
+      }
+      // Clear this bit from every shadow byte.
+      *Ptr &= ~WordValue;
+    }
+  }
+  return WorkingSetSize;
+}
+
+// Scan shadow memory to calculate the number of cache lines being accessed,
+// i.e., the number of non-zero bits indexed by BitIdx in each shadow byte.
+// We also clear the lowest bits (most recent working set snapshot).
+static u32 computeWorkingSizeAndReset(u32 BitIdx) {
+  u32 WorkingSetSize = 0;
+  MemoryMappingLayout MemIter(true/*cache*/);
+  uptr Start, End, Prot;
+  while (MemIter.Next(&Start, &End, nullptr/*offs*/, nullptr/*file*/,
+                      0/*file size*/, &Prot)) {
+    VPrintf(4, "%s: considering %p-%p app=%d shadow=%d prot=%u\n",
+            __FUNCTION__, Start, End, Prot, isAppMem(Start),
+            isShadowMem(Start));
+    if (isShadowMem(Start) && (Prot & MemoryMappingLayout::kProtectionWrite)) {
+      VPrintf(3, "%s: walking %p-%p\n", __FUNCTION__, Start, End);
+      WorkingSetSize += countAndClearShadowValues(BitIdx, Start, End);
+    }
+  }
+  return WorkingSetSize;
+}
+
 void initializeWorkingSet() {
-  // The shadow mapping assumes 64 so this cannot be changed.
-  CHECK(getFlags()->cache_line_size == 64);
+  CHECK(getFlags()->cache_line_size == CacheLineSize);
   registerMemoryFaultHandler();
 }
 
+static u32 getSizeForPrinting(u32 NumOfCachelines, const char *&Unit) {
+  // We need a constant to avoid software divide support:
+  static const u32 KilobyteCachelines = (0x1 << 10) / CacheLineSize;
+  static const u32 MegabyteCachelines = KilobyteCachelines << 10;
+
+  if (NumOfCachelines > 10 * MegabyteCachelines) {
+    Unit = "MB";
+    return NumOfCachelines / MegabyteCachelines;
+  } else if (NumOfCachelines > 10 * KilobyteCachelines) {
+    Unit = "KB";
+    return NumOfCachelines / KilobyteCachelines;
+  } else {
+    Unit = "Bytes";
+    return NumOfCachelines * CacheLineSize;
+  }
+}
+
 int finalizeWorkingSet() {
-  // FIXME NYI: we need to add memory scanning to report the total lines
-  // touched, and later add sampling to get intermediate values.
-  Report("%s is not finished: nothing yet to report\n", SanitizerToolName);
+  // Get the working set size for the entire execution.
+  u32 NumOfCachelines = computeWorkingSizeAndReset(TotalWorkingSetBitIdx);
+  const char *Unit;
+  u32 Size = getSizeForPrinting(NumOfCachelines, Unit);
+  Report(" %s: the total working set size: %u %s (%u cache lines)\n",
+         SanitizerToolName, Size, Unit, NumOfCachelines);
   return 0;
 }
 
diff --git a/lib/esan/working_set_posix.cpp b/lib/esan/working_set_posix.cpp
index f0191766b..0f36c86fe 100644
--- a/lib/esan/working_set_posix.cpp
+++ b/lib/esan/working_set_posix.cpp
@@ -68,10 +68,16 @@ static void reinstateDefaultHandler(int SigNum) {
 // app to handle it just as the app would do without our tool in place.
 static void handleMemoryFault(int SigNum, void *Info, void *Ctx) {
   if (SigNum == SIGSEGV) {
-
-    // TODO: Add shadow memory fault detection and handling.
-
-    if (AppSigAct.sigaction) {
+    // We rely on si_addr being filled in (thus we do not support old kernels).
+    siginfo_t *SigInfo = (siginfo_t *)Info;
+    uptr Addr = (uptr)SigInfo->si_addr;
+    if (isShadowMem(Addr)) {
+      VPrintf(3, "Shadow fault @%p\n", Addr);
+      uptr PageSize = GetPageSizeCached();
+      int Res = internal_mprotect((void *)RoundDownTo(Addr, PageSize),
+                                  PageSize, PROT_READ|PROT_WRITE);
+      CHECK(Res == 0);
+    } else if (AppSigAct.sigaction) {
       // FIXME: For simplicity we ignore app options including its signal stack
       // (we just use ours) and all the delivery flags.
       AppSigAct.sigaction(SigNum, Info, Ctx);