InstrProf: Support profiling dlopen'd shared libraries

Shared objects are hard.  After this commit, we do the right thing when
profiling two separate shared objects that have been dlopen'd with
`RTLD_LOCAL`, when the main executable is *not* being profiled.

This mainly simplifies the writer logic.

  - At initialization, determine the output filename and truncate the
    file.  Depending on whether shared objects can see each other, this
    may happen multiple times.

  - At exit, each executable writes its own profile in append mode.

<rdar://problem/16918688>

llvm-svn: 209053

7 files changed, 148 insertions, 78 deletions

diff --git a/compiler-rt/lib/profile/InstrProfiling.h b/compiler-rt/lib/profile/InstrProfiling.h
index 57ebf3c3a2e..c5dd641fc88 100644
--- a/compiler-rt/lib/profile/InstrProfiling.h
+++ b/compiler-rt/lib/profile/InstrProfiling.h
@@ -83,8 +83,8 @@ void __llvm_profile_set_filename(const char *Name);
 /*! \brief Register to write instrumentation data to file at exit. */
 int __llvm_profile_register_write_file_atexit(void);
 
-/*! \brief Register the write file function for this executable. */
-void __llvm_profile_register_write_file(void);
+/*! \brief Initialize file handling. */
+void __llvm_profile_initialize_file(void);
 
 /*! \brief Get the magic token for the file format. */
 uint64_t __llvm_profile_get_magic(void);
diff --git a/compiler-rt/lib/profile/InstrProfilingFile.c b/compiler-rt/lib/profile/InstrProfilingFile.c
index 5ed32d56267..779a68225c3 100644
--- a/compiler-rt/lib/profile/InstrProfilingFile.c
+++ b/compiler-rt/lib/profile/InstrProfilingFile.c
@@ -53,75 +53,60 @@ static int writeFile(FILE *File) {
   return 0;
 }
 
-typedef struct __llvm_profile_writer {
-  struct __llvm_profile_writer *Next;
-  int (*Data)(FILE *);
-} __llvm_profile_writer;
-
-__attribute__((weak)) __llvm_profile_writer *__llvm_profile_HeadWriter = NULL;
-static __llvm_profile_writer Writer = {NULL, writeFile};
-
-__attribute__((visibility("hidden")))
-void __llvm_profile_register_write_file(void) {
-  static int HasBeenRegistered = 0;
-
-  if (HasBeenRegistered)
-    return;
-
-  HasBeenRegistered = 1;
-  Writer.Next = __llvm_profile_HeadWriter;
-  __llvm_profile_HeadWriter = &Writer;
-}
-
 static int writeFileWithName(const char *OutputName) {
   int RetVal;
   FILE *OutputFile;
   if (!OutputName || !OutputName[0])
     return -1;
-  OutputFile = fopen(OutputName, "w");
+
+  /* Append to the file to support profiling multiple shared objects. */
+  OutputFile = fopen(OutputName, "a");
   if (!OutputFile)
     return -1;
 
-  __llvm_profile_writer *Writer = __llvm_profile_HeadWriter;
-  if (Writer)
-    for (; Writer; Writer = Writer->Next) {
-      RetVal = Writer->Data(OutputFile);
-      if (RetVal != 0)
-        break;
-    }
-  else
-    // Default to calling this executable's writeFile.
-    RetVal = writeFile(OutputFile);
+  RetVal = writeFile(OutputFile);
 
   fclose(OutputFile);
   return RetVal;
 }
 
+__attribute__((weak)) int __llvm_profile_OwnsFilename = 0;
 __attribute__((weak)) const char *__llvm_profile_CurrentFilename = NULL;
-__attribute__((weak)) void __llvm_profile_set_filename(const char *Filename) {
+
+static void setFilename(const char *Filename, int OwnsFilename) {
+  if (__llvm_profile_OwnsFilename)
+    free((char *)__llvm_profile_CurrentFilename);
+
   __llvm_profile_CurrentFilename = Filename;
+  __llvm_profile_OwnsFilename = OwnsFilename;
 }
 
-int getpid(void);
-__attribute__((weak)) int __llvm_profile_write_file(void) {
-  char *AllocatedFilename = NULL;
-  int I, J;
-  int RetVal;
+static void truncateCurrentFile(void) {
+  const char *Filename = __llvm_profile_CurrentFilename;
+  if (!Filename || !Filename[0])
+    return;
 
-#define MAX_PID_SIZE 16
-  char PidChars[MAX_PID_SIZE] = { 0 };
-  int PidLength = 0;
-  int NumPids = 0;
+  /* Truncate the file.  Later we'll reopen and append. */
+  FILE *File = fopen(Filename, "w");
+  if (!File)
+    return;
+  fclose(File);
+}
 
-  /* Get the filename. */
-  const char *Filename = __llvm_profile_CurrentFilename;
-#define UPDATE_FILENAME(NextFilename) \
-  if (!Filename || !Filename[0]) Filename = NextFilename
-  UPDATE_FILENAME(getenv("LLVM_PROFILE_FILE"));
-  UPDATE_FILENAME("default.profraw");
-#undef UPDATE_FILENAME
+static void setDefaultFilename(void) { setFilename("default.profraw", 0); }
+
+int getpid(void);
+static int setFilenameFromEnvironment(void) {
+  const char *Filename = getenv("LLVM_PROFILE_FILE");
+  if (!Filename || !Filename[0])
+    return -1;
 
   /* Check the filename for "%p", which indicates a pid-substitution. */
+#define MAX_PID_SIZE 16
+  char PidChars[MAX_PID_SIZE] = {0};
+  int NumPids = 0;
+  int PidLength = 0;
+  int I;
   for (I = 0; Filename[I]; ++I)
     if (Filename[I] == '%' && Filename[++I] == 'p')
       if (!NumPids++) {
@@ -129,43 +114,74 @@ __attribute__((weak)) int __llvm_profile_write_file(void) {
         if (PidLength <= 0)
           return -1;
       }
-  if (NumPids) {
-    /* Allocate enough space for the substituted filename. */
-    AllocatedFilename = (char*)malloc(I + NumPids*(PidLength - 2) + 1);
-    if (!AllocatedFilename)
-      return -1;
-
-    /* Construct the new filename. */
-    for (I = 0, J = 0; Filename[I]; ++I)
-      if (Filename[I] == '%') {
-        if (Filename[++I] == 'p') {
-          memcpy(AllocatedFilename + J, PidChars, PidLength);
-          J += PidLength;
-        }
-        /* Drop any unknown substitutions. */
-      } else
-        AllocatedFilename[J++] = Filename[I];
-    AllocatedFilename[J] = 0;
-
-    /* Actually use the computed name. */
-    Filename = AllocatedFilename;
+  if (!NumPids) {
+    setFilename(Filename, 0);
+    return 0;
   }
 
-  /* Write the file. */
-  RetVal = writeFileWithName(Filename);
+  /* Allocate enough space for the substituted filename. */
+  char *Allocated = (char*)malloc(I + NumPids*(PidLength - 2) + 1);
+  if (!Allocated)
+    return -1;
 
-  /* Free the filename. */
-  if (AllocatedFilename)
-    free(AllocatedFilename);
+  /* Construct the new filename. */
+  int J;
+  for (I = 0, J = 0; Filename[I]; ++I)
+    if (Filename[I] == '%') {
+      if (Filename[++I] == 'p') {
+        memcpy(Allocated + J, PidChars, PidLength);
+        J += PidLength;
+      }
+      /* Drop any unknown substitutions. */
+    } else
+      Allocated[J++] = Filename[I];
+  Allocated[J] = 0;
 
-  return RetVal;
+  /* Use the computed name. */
+  setFilename(Allocated, 1);
+  return 0;
+}
+
+static void setFilenameAutomatically(void) {
+  if (!setFilenameFromEnvironment())
+    return;
+
+  setDefaultFilename();
+}
+
+__attribute__((visibility("hidden")))
+void __llvm_profile_initialize_file(void) {
+  /* Check if the filename has been initialized. */
+  if (__llvm_profile_CurrentFilename)
+    return;
+
+  /* Detect the filename and truncate. */
+  setFilenameAutomatically();
+  truncateCurrentFile();
+}
+
+__attribute__((visibility("hidden")))
+void __llvm_profile_set_filename(const char *Filename) {
+  setFilename(Filename, 0);
+  truncateCurrentFile();
+}
+
+__attribute__((visibility("hidden")))
+int __llvm_profile_write_file(void) {
+  /* Check the filename. */
+  if (!__llvm_profile_CurrentFilename)
+    return -1;
+
+  /* Write the file. */
+  return writeFileWithName(__llvm_profile_CurrentFilename);
 }
 
 static void writeFileWithoutReturn(void) {
   __llvm_profile_write_file();
 }
 
-__attribute__((weak)) int __llvm_profile_register_write_file_atexit(void) {
+__attribute__((visibility("hidden")))
+int __llvm_profile_register_write_file_atexit(void) {
   static int HasBeenRegistered = 0;
 
   if (HasBeenRegistered)
diff --git a/compiler-rt/lib/profile/InstrProfilingRuntime.cc b/compiler-rt/lib/profile/InstrProfilingRuntime.cc
index 280de834b4b..081ecb29e98 100644
--- a/compiler-rt/lib/profile/InstrProfilingRuntime.cc
+++ b/compiler-rt/lib/profile/InstrProfilingRuntime.cc
@@ -21,7 +21,7 @@ class RegisterRuntime {
 public:
   RegisterRuntime() {
     __llvm_profile_register_write_file_atexit();
-    __llvm_profile_register_write_file();
+    __llvm_profile_initialize_file();
   }
 };
 
diff --git a/compiler-rt/test/profile/Inputs/instrprof-dlopen-func.c b/compiler-rt/test/profile/Inputs/instrprof-dlopen-func.c
new file mode 100644
index 00000000000..f2de3883535
--- /dev/null
+++ b/compiler-rt/test/profile/Inputs/instrprof-dlopen-func.c
@@ -0,0 +1 @@
+void func(int K) { if (K) {} }
diff --git a/compiler-rt/test/profile/Inputs/instrprof-dlopen-func2.c b/compiler-rt/test/profile/Inputs/instrprof-dlopen-func2.c
new file mode 100644
index 00000000000..d4d93dc0b25
--- /dev/null
+++ b/compiler-rt/test/profile/Inputs/instrprof-dlopen-func2.c
@@ -0,0 +1 @@
+void func2(int K) { if (K) {} }
diff --git a/compiler-rt/test/profile/Inputs/instrprof-dlopen-main.c b/compiler-rt/test/profile/Inputs/instrprof-dlopen-main.c
new file mode 100644
index 00000000000..60ee8968360
--- /dev/null
+++ b/compiler-rt/test/profile/Inputs/instrprof-dlopen-main.c
@@ -0,0 +1,18 @@
+#ifdef DLOPEN_FUNC_DIR
+#include <dlfcn.h>
+#else
+void func(int K);
+void func2(int K);
+#endif
+
+int main(int argc, char *argv[]) {
+#ifdef DLOPEN_FUNC_DIR
+  void *f1_handle = dlopen(DLOPEN_FUNC_DIR"/func.shared", DLOPEN_FLAGS);
+  void (*func)(int) = (void (*)(int))dlsym(f1_handle, "func");
+  void *f2_handle = dlopen(DLOPEN_FUNC_DIR"/func2.shared", DLOPEN_FLAGS);
+  void (*func2)(int) = (void (*)(int))dlsym(f2_handle, "func2");
+#endif
+  func(1);
+  func2(0);
+  return 0;
+}
diff --git a/compiler-rt/test/profile/instrprof-dlopen.test b/compiler-rt/test/profile/instrprof-dlopen.test
new file mode 100644
index 00000000000..bba21334f3f
--- /dev/null
+++ b/compiler-rt/test/profile/instrprof-dlopen.test
@@ -0,0 +1,34 @@
+RUN: mkdir -p %t.d
+RUN: %clang_profgen -o %t.d/func.shared -fPIC -shared %S/Inputs/instrprof-dlopen-func.c
+RUN: %clang_profgen -o %t.d/func2.shared -fPIC -shared %S/Inputs/instrprof-dlopen-func2.c
+RUN: %clang -o %t-local -fPIC -DDLOPEN_FUNC_DIR=\"%t.d\" -DDLOPEN_FLAGS=RTLD_LOCAL %S/Inputs/instrprof-dlopen-main.c
+RUN: %clang -o %t-global -fPIC -DDLOPEN_FUNC_DIR=\"%t.d\" -DDLOPEN_FLAGS=RTLD_GLOBAL %S/Inputs/instrprof-dlopen-main.c
+
+RUN: %clang -c -o %t.d/main.o %S/Inputs/instrprof-dlopen-main.c
+RUN: %clang_profgen -o %t-static %S/Inputs/instrprof-dlopen-func.c %S/Inputs/instrprof-dlopen-func2.c %t.d/main.o
+
+RUN: env LLVM_PROFILE_FILE=%t-static.profraw %run %t-static
+RUN: env LLVM_PROFILE_FILE=%t-local.profraw %run %t-local
+RUN: env LLVM_PROFILE_FILE=%t-global.profraw %run %t-global
+
+RUN: llvm-profdata merge -o %t-static.profdata %t-static.profraw
+RUN: llvm-profdata merge -o %t-local.profdata %t-local.profraw
+RUN: llvm-profdata merge -o %t-global.profdata %t-global.profraw
+
+RUN: %clang_profuse=%t-static.profdata -o %t-func.static.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-func.c
+RUN: %clang_profuse=%t-local.profdata -o %t-func.local.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-func.c
+RUN: %clang_profuse=%t-global.profdata -o %t-func.global.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-func.c
+RUN: diff %t-func.static.ll %t-func.local.ll
+RUN: diff %t-func.static.ll %t-func.global.ll
+
+RUN: %clang_profuse=%t-static.profdata -o %t-func2.static.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-func2.c
+RUN: %clang_profuse=%t-local.profdata -o %t-func2.local.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-func2.c
+RUN: %clang_profuse=%t-global.profdata -o %t-func2.global.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-func2.c
+RUN: diff %t-func2.static.ll %t-func2.local.ll
+RUN: diff %t-func2.static.ll %t-func2.global.ll
+
+RUN: %clang_profuse=%t-static.profdata -o %t-main.static.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-main.c
+RUN: %clang_profuse=%t-local.profdata -o %t-main.local.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-main.c
+RUN: %clang_profuse=%t-local.profdata -o %t-main.global.ll -S -emit-llvm %S/Inputs/instrprof-dlopen-main.c
+RUN: diff %t-main.static.ll %t-main.local.ll
+RUN: diff %t-main.static.ll %t-main.global.ll