8 files changed, 42 insertions, 18 deletions
diff --git a/compiler-rt/include/sanitizer/esan_interface.h b/compiler-rt/include/sanitizer/esan_interface.h
index 31e9b9b4df3..4aff8d47bbd 100644
--- a/compiler-rt/include/sanitizer/esan_interface.h
+++ b/compiler-rt/include/sanitizer/esan_interface.h
@@ -39,6 +39,10 @@ extern "C" {
 // data for that point in the run be reported from the tool.
 void COMPILER_RT_WEAK __esan_report();
 
+// This function returns the number of samples that the esan tool has collected
+// to this point.  This is useful for testing.
+unsigned int COMPILER_RT_WEAK __esan_get_sample_count();
+
 #ifdef __cplusplus
 } // extern "C"
 #endif
diff --git a/compiler-rt/lib/esan/esan.cpp b/compiler-rt/lib/esan/esan.cpp
index 3c69b4e91b5..2fb77894d4f 100644
--- a/compiler-rt/lib/esan/esan.cpp
+++ b/compiler-rt/lib/esan/esan.cpp
@@ -259,4 +259,12 @@ void processCompilationUnitExit(void *Ptr) {
   }
 }
 
+unsigned int getSampleCount() {
+  VPrintf(1, "in esan::%s\n", __FUNCTION__);
+  if (__esan_which_tool == ESAN_WorkingSet) {
+    return getSampleCountWorkingSet();
+  }
+  return 0;
+}
+
 } // namespace __esan
diff --git a/compiler-rt/lib/esan/esan.h b/compiler-rt/lib/esan/esan.h
index 371810d5eba..5a0dde62788 100644
--- a/compiler-rt/lib/esan/esan.h
+++ b/compiler-rt/lib/esan/esan.h
@@ -38,6 +38,7 @@ extern bool EsanDuringInit;
 void initializeLibrary(ToolType Tool);
 int finalizeLibrary();
 void reportResults();
+unsigned int getSampleCount();
 // Esan creates the variable per tool per compilation unit at compile time
 // and passes its pointer Ptr to the runtime library.
 void processCompilationUnitInit(void *Ptr);
diff --git a/compiler-rt/lib/esan/esan_interface.cpp b/compiler-rt/lib/esan/esan_interface.cpp
index 8a64d152611..43b3dff86f7 100644
--- a/compiler-rt/lib/esan/esan_interface.cpp
+++ b/compiler-rt/lib/esan/esan_interface.cpp
@@ -115,4 +115,8 @@ extern "C" {
 SANITIZER_INTERFACE_ATTRIBUTE void __esan_report() {
   reportResults();
 }
+
+SANITIZER_INTERFACE_ATTRIBUTE unsigned int __esan_get_sample_count() {
+  return getSampleCount();
+}
 } // extern "C"
diff --git a/compiler-rt/lib/esan/working_set.cpp b/compiler-rt/lib/esan/working_set.cpp
index 3fde5a8b5fc..f39111993c3 100644
--- a/compiler-rt/lib/esan/working_set.cpp
+++ b/compiler-rt/lib/esan/working_set.cpp
@@ -190,6 +190,11 @@ static void takeSample(void *Arg) {
   }
 }
 
+unsigned int getSampleCountWorkingSet()
+{
+  return SnapshotNum;
+}
+
 // Initialization that must be done before any instrumented code is executed.
 void initializeShadowWorkingSet() {
   CHECK(getFlags()->cache_line_size == CacheLineSize);
diff --git a/compiler-rt/lib/esan/working_set.h b/compiler-rt/lib/esan/working_set.h
index 38ff0635d0a..6a976c3f9b2 100644
--- a/compiler-rt/lib/esan/working_set.h
+++ b/compiler-rt/lib/esan/working_set.h
@@ -24,6 +24,7 @@ void initializeWorkingSet();
 void initializeShadowWorkingSet();
 int finalizeWorkingSet();
 void reportWorkingSet();
+unsigned int getSampleCountWorkingSet();
 void processRangeAccessWorkingSet(uptr PC, uptr Addr, SIZE_T Size,
                                   bool IsWrite);
 
diff --git a/compiler-rt/test/esan/TestCases/workingset-midreport.cpp b/compiler-rt/test/esan/TestCases/workingset-midreport.cpp
index 470e33ca36e..2c29cf48ccf 100644
--- a/compiler-rt/test/esan/TestCases/workingset-midreport.cpp
+++ b/compiler-rt/test/esan/TestCases/workingset-midreport.cpp
@@ -17,15 +17,15 @@ const int iters = 6;
 int main(int argc, char **argv) {
   char *buf = (char *)mmap(0, size, PROT_READ | PROT_WRITE,
                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-  // Try to increase the probability that the sideline thread is
-  // scheduled.  Unfortunately we can't do proper synchronization
-  // without some form of annotation or something.
-  sched_yield();
-  // Do enough work to get at least 4 samples.
-  for (int j = 0; j < iters; ++j) {
-    for (int i = 0; i < size; ++i)
-      buf[i] = i;
-    sched_yield();
+  // To avoid flakiness stemming from whether the sideline thread
+  // is scheduled enough on a loaded test machine, we coordinate
+  // with esan itself:
+  if (__esan_get_sample_count) {
+    while (__esan_get_sample_count() < 4) {
+      for (int i = 0; i < size; ++i)
+        buf[i] = i;
+      sched_yield();
+    }
   }
   // Ensure a non-esan build works without ifdefs:
   if (__esan_report) {
diff --git a/compiler-rt/test/esan/TestCases/workingset-samples.cpp b/compiler-rt/test/esan/TestCases/workingset-samples.cpp
index 09efcd44fbf..cf198d2f39e 100644
--- a/compiler-rt/test/esan/TestCases/workingset-samples.cpp
+++ b/compiler-rt/test/esan/TestCases/workingset-samples.cpp
@@ -1,6 +1,7 @@
 // RUN: %clang_esan_wset -O0 %s -o %t 2>&1
 // RUN: %run %t 2>&1 | FileCheck %s
 
+#include <sanitizer/esan_interface.h>
 #include <sched.h>
 #include <stdlib.h>
 #include <string.h>
@@ -12,15 +13,15 @@ const int iters = 6;
 int main(int argc, char **argv) {
   char *buf = (char *)mmap(0, size, PROT_READ | PROT_WRITE,
                            MAP_PRIVATE | MAP_ANONYMOUS, -1, 0);
-  // Try to increase the probability that the sideline thread is
-  // scheduled.  Unfortunately we can't do proper synchronization
-  // without some form of annotation or something.
-  sched_yield();
-  // Do enough work to get at least 4 samples.
-  for (int j = 0; j < iters; ++j) {
-    for (int i = 0; i < size; ++i)
-      buf[i] = i;
-    sched_yield();
+  // To avoid flakiness stemming from whether the sideline thread
+  // is scheduled enough on a loaded test machine, we coordinate
+  // with esan itself:
+  if (__esan_get_sample_count) {
+    while (__esan_get_sample_count() < 4) {
+      for (int i = 0; i < size; ++i)
+        buf[i] = i;
+      sched_yield();
+    }
   }
   munmap(buf, size);
   // We only check for a few samples here to reduce the chance of flakiness.