From 68ea360ed1cd30baa0d11737f7223d2e429262bd Mon Sep 17 00:00:00 2001
From: Martin Pelikan <martin.pelikan@gmail.com>
Date: Wed, 4 Oct 2017 05:12:00 +0000
Subject: [PATCH] [XRay] [compiler-rt] make sure single threaded programs get
 traced too

Summary:
When the XRay user calls the API to finish writing the log, the thread
which is calling the API still hasn't finished and therefore won't get
its trace written.  Add a test for only the main thread to check this.

Reviewers: dberris

Subscribers: llvm-commits

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

llvm-svn: 314875
---
 compiler-rt/lib/xray/xray_fdr_logging.cc      | 10 +++++
 .../xray/TestCases/Linux/fdr-single-thread.cc | 37 +++++++++++++++++++
 2 files changed, 47 insertions(+)
 create mode 100644 compiler-rt/test/xray/TestCases/Linux/fdr-single-thread.cc

diff --git a/compiler-rt/lib/xray/xray_fdr_logging.cc b/compiler-rt/lib/xray/xray_fdr_logging.cc
index 7aebcb224667..bd79893f6186 100644
--- a/compiler-rt/lib/xray/xray_fdr_logging.cc
+++ b/compiler-rt/lib/xray/xray_fdr_logging.cc
@@ -116,6 +116,16 @@ XRayLogFlushStatus fdrLoggingFlush() XRAY_NEVER_INSTRUMENT {
                        reinterpret_cast<char *>(B.Buffer) + B.Size);
     }
   });
+
+  // The buffer for this particular thread would have been finalised after
+  // we've written everything to disk, and we'd lose the thread's trace.
+  auto &TLD = __xray::__xray_fdr_internal::getThreadLocalData();
+  if (TLD.Buffer.Buffer != nullptr) {
+    __xray::__xray_fdr_internal::writeEOBMetadata();
+    auto Start = reinterpret_cast<char *>(TLD.Buffer.Buffer);
+    retryingWriteAll(Fd, Start, Start + TLD.Buffer.Size);
+  }
+
   __sanitizer::atomic_store(&LogFlushStatus,
                             XRayLogFlushStatus::XRAY_LOG_FLUSHED,
                             __sanitizer::memory_order_release);
diff --git a/compiler-rt/test/xray/TestCases/Linux/fdr-single-thread.cc b/compiler-rt/test/xray/TestCases/Linux/fdr-single-thread.cc
new file mode 100644
index 000000000000..30d834e57b0a
--- /dev/null
+++ b/compiler-rt/test/xray/TestCases/Linux/fdr-single-thread.cc
@@ -0,0 +1,37 @@
+// RUN: %clangxx_xray -g -std=c++11 %s -o %t
+// RUN: XRAY_OPTIONS=XRAY_OPTIONS="verbosity=1 patch_premain=true \
+// RUN:   xray_naive_log=false xray_fdr_log=true \
+// RUN:   xray_fdr_log_func_duration_threshold_us=0 \
+// RUN:   xray_logfile_base=fdr-logging-1thr-" %run %t 2>&1
+// RUN: %llvm_xray convert --output-format=yaml --symbolize --instr_map=%t \
+// RUN:   "`ls fdr-logging-1thr-* | head -n1`" | FileCheck %s
+// RUN: rm fdr-logging-1thr-*
+//
+// REQUIRES: x86_64-linux
+
+#include "xray/xray_log_interface.h"
+#include <cassert>
+
+constexpr auto kBufferSize = 16384;
+constexpr auto kBufferMax = 10;
+
+[[clang::xray_always_instrument]] void __attribute__((noinline)) fn() { }
+
+int main(int argc, char *argv[]) {
+  using namespace __xray;
+  FDRLoggingOptions Opts;
+
+  auto status = __xray_log_init(kBufferSize, kBufferMax, &Opts, sizeof(Opts));
+  assert(status == XRayLogInitStatus::XRAY_LOG_INITIALIZED);
+
+  __xray_patch();
+  fn();
+  __xray_unpatch();
+  assert(__xray_log_finalize() == XRAY_LOG_FINALIZED);
+  assert(__xray_log_flushLog() == XRAY_LOG_FLUSHED);
+  return 0;
+}
+
+// CHECK: records:
+// CHECK-NEXT: - { type: 0, func-id: [[FID1:[0-9]+]], function: {{.*fn.*}}, cpu: {{.*}}, thread: [[THREAD1:[0-9]+]], kind: function-enter, tsc: {{[0-9]+}} }
+// CHECK-NEXT: - { type: 0, func-id: [[FID1:[0-9]+]], function: {{.*fn.*}}, cpu: {{.*}}, thread: [[THREAD1:[0-9]+]], kind: function-exit, tsc: {{[0-9]+}} }