[profile] Add support for the exit-on-signal '%Nx' specifier

The exit-on-signal specifier tells the runtime to handle a signal by writing
out a profile and then exiting the program. It can be specified in the
LLVM_PROFILE_FILE environment variable or by calling
__llvm_profile_set_filename.

Here is how you might force a program to exit when it's sent SIGTERM (sig 15):

  LLVM_PROFILE_FILE="default.profraw%15x" <program>

When an instance of <program> is sent SIGTERM, it will write out a profile to
"default.profraw" (the specifier is stripped out of the filename) and exit with
return code 0. It's possible to use multiple exit-on-signal specifiers (up to
16).

rdar://problem/24098975

Author: vedantk

lib/profile/InstrProfilingFile.c

@@ -14,6 +14,7 @@
 #include <stdio.h>
 #include <stdlib.h>
 #include <string.h>
+#include <signal.h>
 #ifdef _MSC_VER
 /* For _alloca. */
 #include <malloc.h>
@@ -59,6 +60,7 @@ static const char *getPNSStr(ProfileNameSpecifier PNS) {
 }
 
 #define MAX_PID_SIZE 16
+#define MAX_SIGNAL_HANDLERS 16
 /* Data structure holding the result of parsed filename pattern. */
 typedef struct lprofFilename {
   /* File name string possibly with %p or %h specifiers. */
@@ -79,11 +81,13 @@ typedef struct lprofFilename {
    * 2 profile data files. %1m is equivalent to %m. Also %m specifier
    * can only appear once at the end of the name pattern. */
   unsigned MergePoolSize;
+  char ExitOnSignals[MAX_SIGNAL_HANDLERS];
+  unsigned NumExitSignals;
   ProfileNameSpecifier PNS;
 } lprofFilename;
 
-COMPILER_RT_WEAK lprofFilename lprofCurFilename = {0, 0, 0, {0}, {0},
-                                                   0, 0, 0, PNS_unknown};
+COMPILER_RT_WEAK lprofFilename lprofCurFilename = {
+    0, 0, 0, {0}, {0}, 0, 0, 0, {0}, 0, PNS_unknown};
 
 int getpid(void);
 static int getCurFilenameLength();
@@ -252,6 +256,26 @@ static void truncateCurrentFile(void) {
   fclose(File);
 }
 
+static void exitSignalHandler(int sig) {
+  (void)sig;
+  exit(0);
+}
+
+static void installExitSignalHandlers(void) {
+  unsigned I;
+  struct sigaction sigact;
+  int err;
+  for (I = 0; I < lprofCurFilename.NumExitSignals; ++I) {
+    memset(&sigact, 0, sizeof(sigact));
+    sigact.sa_handler = exitSignalHandler;
+    err = sigaction(lprofCurFilename.ExitOnSignals[I], &sigact, NULL);
+    if (err)
+      PROF_WARN(
+          "Unable to install an exit signal handler for %d (errno = %d).\n",
+          lprofCurFilename.ExitOnSignals[I], err);
+  }
+}
+
 static const char *DefaultProfileName = "default.profraw";
 static void resetFilenameToDefault(void) {
   if (lprofCurFilename.FilenamePat && lprofCurFilename.OwnsFilenamePat) {
@@ -262,14 +286,23 @@ static void resetFilenameToDefault(void) {
   lprofCurFilename.PNS = PNS_default;
 }
 
+static int isNonZeroDigit(char C) { return C >= '1' && C <= '9'; }
+
 static int containsMergeSpecifier(const char *FilenamePat, int I) {
   return (FilenamePat[I] == 'm' ||
-          (FilenamePat[I] >= '1' && FilenamePat[I] <= '9' &&
+          (isNonZeroDigit(FilenamePat[I]) &&
            /* If FilenamePat[I] is not '\0', the next byte is guaranteed
             * to be in-bound as the string is null terminated. */
            FilenamePat[I + 1] == 'm'));
 }
 
+static int containsExitOnSignalSpecifier(const char *FilenamePat, int I) {
+  if (!isNonZeroDigit(FilenamePat[I]))
+    return 0;
+  return (FilenamePat[I + 1] == 'x') ||
+         (isNonZeroDigit(FilenamePat[I + 1]) && FilenamePat[I + 2] == 'x');
+}
+
 /* Parses the pattern string \p FilenamePat and stores the result to
  * lprofcurFilename structure. */
 static int parseFilenamePattern(const char *FilenamePat,
@@ -278,6 +311,7 @@ static int parseFilenamePattern(const char *FilenamePat,
   char *PidChars = &lprofCurFilename.PidChars[0];
   char *Hostname = &lprofCurFilename.Hostname[0];
   int MergingEnabled = 0;
+  char SignalNo;
 
   /* Clean up cached prefix.  */
   if (lprofCurFilename.ProfilePathPrefix)
@@ -327,6 +361,22 @@ static int parseFilenamePattern(const char *FilenamePat,
           lprofCurFilename.MergePoolSize = FilenamePat[I] - '0';
           I++; /* advance to 'm' */
         }
+      } else if (containsExitOnSignalSpecifier(FilenamePat, I)) {
+        if (lprofCurFilename.NumExitSignals == MAX_SIGNAL_HANDLERS) {
+          PROF_WARN("%%x specifier has been specified too many times in %s.\n",
+                    FilenamePat);
+          return -1;
+        }
+        /* Grab the signal number. */
+        SignalNo = FilenamePat[I] - '0';
+        I++; /* advance to either another digit, or 'x' */
+        if (FilenamePat[I] != 'x') {
+          SignalNo = (SignalNo * 10) + (FilenamePat[I] - '0');
+          I++; /* advance to 'x' */
+        }
+        lprofCurFilename.ExitOnSignals[lprofCurFilename.NumExitSignals] =
+            SignalNo;
+        ++lprofCurFilename.NumExitSignals;
       }
     }
 
@@ -370,6 +420,7 @@ static void parseAndSetFilename(const char *FilenamePat,
   }
 
   truncateCurrentFile();
+  installExitSignalHandlers();
 }
 
 /* Return buffer length that is required to store the current profile
@@ -378,18 +429,24 @@ static void parseAndSetFilename(const char *FilenamePat,
 #define SIGLEN 24
 static int getCurFilenameLength() {
   int Len;
+  unsigned I;
   if (!lprofCurFilename.FilenamePat || !lprofCurFilename.FilenamePat[0])
     return 0;
 
   if (!(lprofCurFilename.NumPids || lprofCurFilename.NumHosts ||
-        lprofCurFilename.MergePoolSize))
+        lprofCurFilename.MergePoolSize || lprofCurFilename.NumExitSignals))
     return strlen(lprofCurFilename.FilenamePat);
 
   Len = strlen(lprofCurFilename.FilenamePat) +
         lprofCurFilename.NumPids * (strlen(lprofCurFilename.PidChars) - 2) +
         lprofCurFilename.NumHosts * (strlen(lprofCurFilename.Hostname) - 2);
   if (lprofCurFilename.MergePoolSize)
     Len += SIGLEN;
+  for (I = 0; I < lprofCurFilename.NumExitSignals; ++I) {
+    Len -= 3; /* Drop the '%', signal number, and the 'x'. */
+    if (lprofCurFilename.ExitOnSignals[I] >= 10)
+      --Len; /* Drop the second digit of the signal number. */
+  }
   return Len;
 }
 
@@ -405,7 +462,7 @@ static const char *getCurFilename(char *FilenameBuf) {
     return 0;
 
   if (!(lprofCurFilename.NumPids || lprofCurFilename.NumHosts ||
-        lprofCurFilename.MergePoolSize))
+        lprofCurFilename.MergePoolSize || lprofCurFilename.NumExitSignals))
     return lprofCurFilename.FilenamePat;
 
   PidLength = strlen(lprofCurFilename.PidChars);
@@ -431,6 +488,9 @@ static const char *getCurFilename(char *FilenameBuf) {
         J += S;
         if (FilenamePat[I] != 'm')
           I++;
+      } else if (containsExitOnSignalSpecifier(FilenamePat, I)) {
+        while (FilenamePat[I] != 'x')
+          ++I;
       }
       /* Drop any unknown substitutions. */
     } else

test/profile/instrprof-exit-on-signal.c

@@ -0,0 +1,16 @@
+// RUN: %clang_profgen -o %t %s
+// RUN: %run LLVM_PROFILE_FILE="%15x%t.profraw" %t
+// RUN: llvm-profdata show --all-functions %t.profraw | FileCheck %s
+// CHECK: Total functions: 1
+
+#include <signal.h>
+#include <unistd.h>
+
+int main() {
+  kill(getpid(), SIGTERM);
+  
+  while (1) {
+    /* loop forever */
+  }
+  return 1;
+}