Release pages to OS when setting 0 label

[stephan.yichao.zhao]

Summary

This is a follow up patch of https://reviews.llvm.org/D88755.

When set 0 label for an address range, we can release pages within the
corresponding shadow address range to OS, and set only addresses outside
the pages to be 0.

[morehouse]

compiler-rt/lib/dfsan/dfsan.cpp

277–280

Nit: Let's rename this something more descriptive. e.g., WriteShadowIfDifferent.

278

Every call site seems to cast a uptr to dfsan_label *. Let's move the cast into this function and rename the parameter something more descriptive (e.g., shadow_addr).

307

Both MSan and ASan do the release only when size is larger than 16 pages. Should we have similar logic here?

Maybe it's less important to avoid the system call in this case for DFSan... ASan/MSan can use memset to clear a small number of pages, which can use vector instructions for efficiency. But DFSan has to use a loop since shadow labels take up 2 bytes. So maybe the threshold where madvise becomes faster is lower for DFSan.

@eugenis WDYT?

[eugenis]

compiler-rt/lib/dfsan/dfsan.cpp

307

If you have a specific application in mind, I suggest directly measuring run time and process RSS. Otherwise, the 16 page threshold makes sense to me, AFAIR it had the best trade-off for ASan.

[stephan.yichao.zhao]

compiler-rt/lib/dfsan/dfsan.cpp

307

It seems that 32k is the boundary between calls from malloc/free and utility zero-out. Added comments.

[stephan.yichao.zhao]

addressed comments

[morehouse]

compiler-rt/lib/dfsan/dfsan.cpp

318

Since kNumPagesThreshold is only used here, let's define it here instead.

[stephan.yichao.zhao]

moved constant definition.

compiler-rt/lib/dfsan/dfsan.cpp

dfsan_label label =
       atomic_fetch_add(&__dfsan_last_label, 1, memory_order_relaxed) + 1;
   dfsan_check_label(label);
   __dfsan_label_info[label].l1 = __dfsan_label_info[label].l2 = 0;
   __dfsan_label_info[label].desc = desc;
   __dfsan_label_info[label].userdata = userdata;
   return label;
 }
 
-extern "C" SANITIZER_INTERFACE_ATTRIBUTE
-void __dfsan_set_label(dfsan_label label, void *addr, uptr size) {
-  for (dfsan_label *labelp = shadow_for(addr); size != 0; --size, ++labelp) {
+static void WriteShadowIfDifferent(dfsan_label label, uptr shadow_addr,
+                                   uptr size) {
+  dfsan_label *labelp = (dfsan_label *)shadow_addr;
+  for (; size != 0; --size, ++labelp) {
     // Don't write the label if it is already the value we need it to be.
     // In a program where most addresses are not labeled, it is common that
     // a page of shadow memory is entirely zeroed.  The Linux copy-on-write
     // implementation will share all of the zeroed pages, making a copy of a
     // page when any value is written.  The un-sharing will happen even if
     // the value written does not change the value in memory.  Avoiding the
     // write when both |label| and |*labelp| are zero dramatically reduces
     // the amount of real memory used by large programs.
     if (label == *labelp)
       continue;
 
     *labelp = label;
   }
 }
 
+extern "C" SANITIZER_INTERFACE_ATTRIBUTE void __dfsan_set_label(
+    dfsan_label label, void *addr, uptr size) {
+  const uptr beg_shadow_addr = (uptr)__dfsan::shadow_for(addr);
+
+  if (0 != label) {
+    WriteShadowIfDifferent(label, beg_shadow_addr, size);
+    return;
+  }
+
+  // If label is 0, releases the pages within the shadow address range, and sets
+  // the shadow addresses not on the pages to be 0.
+  const void *end_addr = (void *)((uptr)addr + size);
+  const uptr end_shadow_addr = (uptr)__dfsan::shadow_for(end_addr);
+  const uptr page_size = GetPageSizeCached();
+  const uptr beg_aligned = RoundUpTo(beg_shadow_addr, page_size);
+  const uptr end_aligned = RoundDownTo(end_shadow_addr, page_size);
+
+  // dfsan_set_label can be called from the following cases
+  // 1) mapped ranges by new/delete and malloc/free. This case has shadow memory
+  // size > 100k, and happens less frequently.
+  // 2) zero-filling internal data structures by utility libraries. This case
+  // has shadow memory size < 32k, and happens more often.
+  // Set kNumPagesThreshold to be 8 to avoid releasing small pages.
+  const int kNumPagesThreshold = 8;
+  if (beg_aligned + kNumPagesThreshold * page_size >= end_aligned)
+    return WriteShadowIfDifferent(label, beg_shadow_addr, size);
+
+  WriteShadowIfDifferent(label, beg_shadow_addr, beg_aligned - beg_shadow_addr);
+  ReleaseMemoryPagesToOS(beg_aligned, end_aligned);
+  WriteShadowIfDifferent(label, end_aligned, end_shadow_addr - end_aligned);
+}
+
 SANITIZER_INTERFACE_ATTRIBUTE
 void dfsan_set_label(dfsan_label label, void *addr, uptr size) {
   __dfsan_set_label(label, addr, size);
 }
 
 SANITIZER_INTERFACE_ATTRIBUTE
 void dfsan_add_label(dfsan_label label, void *addr, uptr size) {
   for (dfsan_label *labelp = shadow_for(addr); size != 0; --size, ++labelp)

compiler-rt/lib/dfsan/dfsan_interceptors.cpp

 #include "sanitizer_common/sanitizer_common.h"
 
 using namespace __sanitizer;
 
 namespace {
 
 bool interceptors_initialized;
 
-void ReleaseShadowMemoryPagesToOS(void *addr, SIZE_T length) {
-  uptr beg_shadow_addr = (uptr)__dfsan::shadow_for(addr);
-  void *end_addr =
-      (void *)((uptr)addr + RoundUpTo(length, GetPageSizeCached()));
-  uptr end_shadow_addr = (uptr)__dfsan::shadow_for(end_addr);
-  ReleaseMemoryPagesToOS(beg_shadow_addr, end_shadow_addr);
-}
-
 }  // namespace
 
 INTERCEPTOR(void *, mmap, void *addr, SIZE_T length, int prot, int flags,
             int fd, OFF_T offset) {
   void *res;
 
   // interceptors_initialized is set to true during preinit_array, when we're
   // single-threaded.  So we don't need to worry about accessing it atomically.
   if (!interceptors_initialized)
     res = (void *)syscall(__NR_mmap, addr, length, prot, flags, fd, offset);
   else
     res = REAL(mmap)(addr, length, prot, flags, fd, offset);
 
   if (res != (void *)-1)
-    ReleaseShadowMemoryPagesToOS(res, length);
+    dfsan_set_label(0, res, RoundUpTo(length, GetPageSizeCached()));
   return res;
 }
 
 INTERCEPTOR(void *, mmap64, void *addr, SIZE_T length, int prot, int flags,
             int fd, OFF64_T offset) {
   void *res = REAL(mmap64)(addr, length, prot, flags, fd, offset);
   if (res != (void *)-1)
-    ReleaseShadowMemoryPagesToOS(res, length);
+    dfsan_set_label(0, res, RoundUpTo(length, GetPageSizeCached()));
   return res;
 }
 
 INTERCEPTOR(int, munmap, void *addr, SIZE_T length) {
   int res = REAL(munmap)(addr, length);
-  if (res != -1) {
-    ReleaseShadowMemoryPagesToOS(addr, length);
-  }
+  if (res != -1)
+    dfsan_set_label(0, addr, RoundUpTo(length, GetPageSizeCached()));
   return res;
 }
 
 namespace __dfsan {
 void InitializeInterceptors() {
   CHECK(!interceptors_initialized);
 
   INTERCEPT_FUNCTION(mmap);
   INTERCEPT_FUNCTION(mmap64);
   INTERCEPT_FUNCTION(munmap);
 
   interceptors_initialized = true;
 }
 }  // namespace __dfsan