[PGO] Suspend SIGKILL for PR_SET_PDEATHSIG in profile-write

Summary:
We found a nondeterministic behavior when doing online profile merging
for multi-process applications. The application forks a sub-process and
sub-process sets to get SIGKILL when the parent process exits,

The first process gets the lock, and dumps the profile. The second one
will mmap the file, do the merge and write out the file. Note that before
the merged write, we truncate the profile.

Depending on the timing, the child process might be terminated
abnormally when the parent exits first. If this happens:
(1) before the truncation, we will get the profile for the main process
(2) after the truncation, and before write-out the profile,  we will get
0 size profile.
(3) after the merged write, we get merged profile.

This patch temporarily suspend the SIGKILL for PR_SET_PDEATHSIG
before profile-write and restore it after the write.

This patch only applies to Linux system.

Reviewers: davidxl

Reviewed By: davidxl

Subscribers: xur, llvm-commits

Differential Revision: https://reviews.llvm.org/D29954

llvm-svn: 295364

4 files changed, 79 insertions, 0 deletions

diff --git a/compiler-rt/lib/profile/InstrProfilingFile.c b/compiler-rt/lib/profile/InstrProfilingFile.c
index cd3590e12f5..dfcbe52d7e4 100644
--- a/compiler-rt/lib/profile/InstrProfilingFile.c
+++ b/compiler-rt/lib/profile/InstrProfilingFile.c
@@ -530,6 +530,7 @@ int __llvm_profile_write_file(void) {
   int rc, Length;
   const char *Filename;
   char *FilenameBuf;
+  int PDeathSig = 0;
 
   if (lprofProfileDumped()) {
     PROF_NOTE("Profile data not written to file: %s.\n", 
@@ -556,10 +557,18 @@ int __llvm_profile_write_file(void) {
     return -1;
   }
 
+  // Temporarily suspend getting SIGKILL when the parent exits.
+  PDeathSig = lprofSuspendSigKill();
+
   /* Write profile data to the file. */
   rc = writeFile(Filename);
   if (rc)
     PROF_ERR("Failed to write file \"%s\": %s\n", Filename, strerror(errno));
+
+  // Restore SIGKILL.
+  if (PDeathSig == 1)
+    lprofRestoreSigKill();
+
   return rc;
 }
 
diff --git a/compiler-rt/lib/profile/InstrProfilingUtil.c b/compiler-rt/lib/profile/InstrProfilingUtil.c
index 321c7192cc6..d1160411f40 100644
--- a/compiler-rt/lib/profile/InstrProfilingUtil.c
+++ b/compiler-rt/lib/profile/InstrProfilingUtil.c
@@ -29,6 +29,11 @@
 #include <stdlib.h>
 #include <string.h>
 
+#if defined(__linux__)
+#include <signal.h>
+#include <sys/prctl.h>
+#endif
+
 COMPILER_RT_VISIBILITY
 void __llvm_profile_recursive_mkdir(char *path) {
   int i;
@@ -219,3 +224,24 @@ COMPILER_RT_VISIBILITY const char *lprofFindLastDirSeparator(const char *Path) {
 #endif
   return Sep;
 }
+
+COMPILER_RT_VISIBILITY int lprofSuspendSigKill() {
+#if defined(__linux__)
+  int PDeachSig = 0;
+  /* Temporarily suspend getting SIGKILL upon exit of the parent process. */
+  if (prctl(PR_GET_PDEATHSIG, &PDeachSig) == 0 && PDeachSig == SIGKILL) {
+    fprintf(stderr, "set\n");
+    prctl(PR_SET_PDEATHSIG, 0);
+  }
+  return (PDeachSig == SIGKILL);
+#else
+  return 0;
+#endif
+}
+
+COMPILER_RT_VISIBILITY void lprofRestoreSigKill() {
+#if defined(__linux__)
+  fprintf(stderr, "restore \n");
+  prctl(PR_SET_PDEATHSIG, SIGKILL);
+#endif
+}
diff --git a/compiler-rt/lib/profile/InstrProfilingUtil.h b/compiler-rt/lib/profile/InstrProfilingUtil.h
index a80fde77e16..9698599606e 100644
--- a/compiler-rt/lib/profile/InstrProfilingUtil.h
+++ b/compiler-rt/lib/profile/InstrProfilingUtil.h
@@ -51,4 +51,12 @@ int lprofGetHostName(char *Name, int Len);
 unsigned lprofBoolCmpXchg(void **Ptr, void *OldV, void *NewV);
 void *lprofPtrFetchAdd(void **Mem, long ByteIncr);
 
+/* Temporarily suspend SIGKILL. Return value of 1 means a restore is needed.
+ * Other return values mean no restore is needed.
+ */
+int lprofSuspendSigKill();
+
+/* Restore previously suspended SIGKILL. */
+void lprofRestoreSigKill();
+
 #endif /* PROFILE_INSTRPROFILINGUTIL_H */
diff --git a/compiler-rt/test/profile/Linux/prctl.c b/compiler-rt/test/profile/Linux/prctl.c
new file mode 100644
index 00000000000..43baf65439e
--- /dev/null
+++ b/compiler-rt/test/profile/Linux/prctl.c
@@ -0,0 +1,36 @@
+// RUN: %clang_pgogen -O2 -o %t %s
+// RUN: rm -rf default_*.profraw
+// RUN: %run %t && sleep 1
+// RUN: llvm-profdata show default_*.profraw 2>&1 | FileCheck %s
+
+#include <signal.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <sys/prctl.h>
+#include <unistd.h>
+
+#define FAKE_COUNT_SZ 2000000
+/* fake counts to increse the profile size. */
+unsigned long long __attribute__((section("__llvm_prf_cnts")))
+counts[FAKE_COUNT_SZ];
+
+int main(int argc, char **argv) {
+  pid_t pid = fork();
+  if (pid == 0) {
+    int i;
+    int sum = 0;
+    /* child process: sleep 500us and get to runtime before the
+     * main process exits. */
+    prctl(PR_SET_PDEATHSIG, SIGKILL);
+    usleep(500);
+    for (i = 0; i < 5000; ++i)
+      sum += i * i * i;
+    printf("child process (%d): sum=%d\n", getpid(), sum);
+  } else if (pid > 0) {
+    /* parent process: sleep 100us to get into profile runtime first. */
+    usleep(100);
+  }
+  return 0;
+}
+
+// CHECK-NOT: Empty raw profile file