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

[bruening]

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.

[aizatsky]

lib/esan/working_set.cpp

82

This function changes shadow value. Should be at least reflected in the name.

Plus "processShadowRegion" is not very telling.

108

Same.

E.g. "computeWorkingSizeAndClearRecent".

134

symbolic constant for 64 please.

136

static

lib/esan/working_set_posix.cpp

69–78

Should this be working_set_linux.cpp then?

[filcab]

lib/esan/working_set.cpp

88

RoundDown? We'll be looking at some shadow bytes before ShadowStart. Won't that be a problem?

92

sizeof(u32)

93

ByteValue has only one bit set. No need to test for equality.

133

We don't really care, since we only do this at finalize time (and we're printing stuff anyway), no?

135

No sense, again. Just make a global kDefaultCacheLineSize, and document next to it that we only support that size.

lib/esan/working_set_posix.cpp

69–78

or at least the _posix variant.

test/esan/TestCases/workingset-simple.cpp

19

What do you mean?

[bruening]

Address reviewer requests.

lib/esan/working_set.cpp

88

No, this routine may be used for handling smaller regions in the future (e.g., just the heap) where it is convenient to not require precision or alignment from the caller.

133

There are linker errors for missing symbols like __cxa_guard_acquire which would require adding libraries these sanitizer libs try to avoid.

lib/esan/working_set_posix.cpp

69–78

POSIX specifies that si_addr be filled in for SIGSEGV. It's older Linux that fails to do so (we're talking really old: 10+ years IIRC).

[filcab]

LGTM

lib/esan/working_set.cpp

35

Should this be in the header?

82

We should probably mention it might count and clear a few bytes more than [ShadowStart,ShadowEnd)

[bruening]

lib/esan/working_set.cpp

35

We can move it to the header once another file needs it.

82

That is what the word-align comment implies.

[vitalybuka]

lib/esan/working_set.cpp

87

Is this the same as?
u32 WordValue = 0x01010101 << BitIdx;

[eugenis]

This revision depends on r270650, which broke the bots. Please fix that first.
http://lab.llvm.org:8011/builders/sanitizer-x86_64-linux/builds/24123/steps/test%20standalone%20compiler-rt/logs/stdio

[bruening]

The bot is now green.

lib/esan/esan.cpp

static void initializeShadow() {
 
   VPrintf(1, "Shadow scale=%d offset=%p\n", Mapping.Scale, Mapping.Offset);
 
   uptr ShadowStart, ShadowEnd;
   for (int i = 0; getShadowRegion(i, &ShadowStart, &ShadowEnd); ++i) {
     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,
+    } else {
+      Map = (uptr)MmapFixedNoReserve(ShadowStart, ShadowEnd - ShadowStart,
+                                     "shadow");
+    }
     if (Map != ShadowStart) {
       Printf("FATAL: EfficiencySanitizer failed to map its shadow memory.\n");
       Die();
     }
 
     if (common_flags()->no_huge_pages_for_shadow)
       NoHugePagesInRegion(ShadowStart, ShadowEnd - ShadowStart);
     if (common_flags()->use_madv_dontdump)

lib/esan/working_set.cpp

 //
 // This file contains working-set-specific code.
 //===----------------------------------------------------------------------===//
 
 #include "working_set.h"
 #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
 //   cache line has ever been accessed.
 // - The lowest bit of each shadow byte indicates whether the corresponding
 //   cache line was accessed since the last sample.
 // - The other bits can be used either for a single working set snapshot
 //   between two consecutive samples, or an aggregate working set snapshot
 //   over multiple sample periods (future work).
 // We live with races in accessing each shadow byte.
 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;
 static const byte ShadowAccessedVal =
   (1 << TotalWorkingSetBitIdx) | (1 << CurWorkingSetBitIdx);
 
 void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
                                   bool IsWrite) {
void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
   while (I < NumLines) {
     if ((*Shadow & ShadowAccessedVal) != ShadowAccessedVal)
       *Shadow |= ShadowAccessedVal;
     ++Shadow;
     ++I;
   }
 }
 
+// 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);
   registerShadowFault();
 }
 
+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;
 }
 
 } // namespace __esan

lib/esan/working_set_posix.cpp

static void reinstateDefaultHandler(int SigNum) {
   SigAct.sigaction = (void (*)(int, void*, void*)) SIG_DFL;
   int Res = internal_sigaction(SigNum, &SigAct, nullptr);
   CHECK(Res == 0);
   VPrintf(1, "Unregistered for %d handler\n", SigNum);
 }
 
 static void handleShadowFault(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);
     } else {
       // Crash instead of spinning with infinite faults.
       reinstateDefaultHandler(SigNum);
     }
   } else

test/esan/TestCases/workingset-fault.cpp

int main(int argc, char **argv) {
   assert(SigAct.sa_sigaction == SigactionHandler);
 
   return 0;
 }
 // CHECK:      Handling SIGSEGV for signal
 // CHECK-NEXT: Past longjmp for signal
 // CHECK-NEXT: Handling SIGSEGV for sigaction
 // CHECK-NEXT: Past longjmp for sigaction
+// CHECK:      {{.*}} EfficiencySanitizer: the total working set size: {{[0-9][0-9][0-9]}} Bytes ({{[0-9][0-9]}} cache lines)

test/esan/TestCases/workingset-memset.cpp

int main(int argc, char **argv) {
   const int size = 128*1024*1024;
   char *p = (char *)mmap(0, size, PROT_READ | PROT_WRITE,
                          MAP_ANON | MAP_PRIVATE, -1, 0);
   // Test the slowpath at different cache line boundaries.
   for (int i = 0; i < 630; i++)
     memset((char *)p + 63*i, i, 63*i);
   munmap(p, size);
   return 0;
-  // FIXME: once the memory scan and size report is in place add it here.
-  // CHECK: {{.*}}EfficiencySanitizer is not finished: nothing yet to report
+  // CHECK: {{.*}} EfficiencySanitizer: the total working set size: 77 KB (12{{[0-9]+}} cache lines)
 }

test/esan/TestCases/workingset-simple.cpp

+// RUN: %clang_esan_wset -O0 %s -o %t 2>&1
+// RUN: %run %t 2>&1 | FileCheck %s
+
+#include <stdlib.h>
+#include <string.h>
+#include <sys/mman.h>
+#include <assert.h>
+
+const int size = 0x1 << 25; // 523288 cache lines
+const int line_size = 64;
+
+int main(int argc, char **argv) {
+  char *bufA = (char *)malloc(sizeof(char) * line_size);
+  char bufB[64];
+  char *bufC = (char *)mmap(0, size, PROT_READ | PROT_WRITE,
+                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
+  bufA[0] = 0;
+  // This additional access to the same line should not increase the line
+  // count: but it's difficult to make a non-flaky test that measures the
+  // lines down to the ones digit so right now we're not really testing that.
+  // If we add a heap-only mode we may be able to be more precise.
+  bufA[1] = 0;
+  bufB[33] = 1;
+  for (int i = 0; i < size; i += line_size)
+    bufC[i] = 0;
+  free(bufA);
+  munmap(bufC, 0x4000);
+  // CHECK: {{.*}} EfficiencySanitizer: the total working set size: 32 MB (524{{[0-9][0-9][0-9]}} cache lines)
+  return 0;
+}