Allow extra params on persistent trace collection

Differential Revision: D36364989

fbshipit-source-id: a2c166a12189e0785aaac45ee5f60606c1e6b227

cpp/mmapbuf/writer/MmapBufferTraceWriter.cpp

@@ -26,6 +26,8 @@
 #include <fstream>
 #include <memory>
 #include <stdexcept>
+#include <utility>
+#include <vector>
 
 #include <profilo/entries/EntryType.h>
 #include <profilo/logger/buffer/RingBuffer.h>
@@ -208,14 +210,25 @@ void MmapBufferTraceWriter::nativeWriteTrace(
     const std::string& trace_folder,
     const std::string& trace_prefix,
     int32_t trace_flags,
-    fbjni::alias_ref<JNativeTraceWriterCallbacks> callbacks) {
+    fbjni::alias_ref<JNativeTraceWriterCallbacks> callbacks,
+    fbjni::alias_ref<fbjni::JArrayClass<jstring>> extraAnnotations) {
+  std::vector<std::pair<std::string, std::string>> parsedAnnotations;
+  auto extrasSize = extraAnnotations->size();
+  for (size_t i = 0; i < extrasSize; i += 2) {
+    parsedAnnotations.emplace_back(
+        extraAnnotations->getElement(i)->toStdString(),
+        i + 1 < extrasSize ? extraAnnotations->getElement(i + 1)->toStdString()
+                           : "null");
+  }
+
   writeTrace(
       type,
       persistent,
       trace_folder,
       trace_prefix,
       trace_flags,
-      std::make_shared<NativeTraceWriterCallbacksProxy>(callbacks));
+      std::make_shared<NativeTraceWriterCallbacksProxy>(callbacks),
+      parsedAnnotations);
 }
 
 void MmapBufferTraceWriter::writeTrace(
@@ -225,6 +238,7 @@ void MmapBufferTraceWriter::writeTrace(
     const std::string& trace_prefix,
     int32_t trace_flags,
     std::shared_ptr<TraceCallbacks> callbacks,
+    const std::vector<std::pair<std::string, std::string>>& extraAnnotations,
     uint64_t timestamp) {
   if (bufferMapHolder_.get() == nullptr) {
     throw std::runtime_error(
@@ -301,6 +315,10 @@ void MmapBufferTraceWriter::writeTrace(
     loggerWriteQplTriggerAnnotation(
         logger, qpl_marker_id, "collection_method", "persistent", timestamp);
   }
+  for (auto extra : extraAnnotations) {
+    loggerWriteQplTriggerAnnotation(
+        logger, qpl_marker_id, extra.first, extra.second, timestamp);
+  }
 
   const char* mapsFilePath = mapBufferPrefix->header.memoryMapsFilePath;
   if (mapsFilePath[0] != '\0') {

cpp/mmapbuf/writer/MmapBufferTraceWriter.h

@@ -54,7 +54,8 @@ class MmapBufferTraceWriter : public fbjni::HybridClass<MmapBufferTraceWriter> {
       const std::string& trace_folder,
       const std::string& trace_prefix,
       int32_t trace_flags,
-      fbjni::alias_ref<JNativeTraceWriterCallbacks> callbacks);
+      fbjni::alias_ref<JNativeTraceWriterCallbacks> callbacks,
+      fbjni::alias_ref<fbjni::JArrayClass<jstring>> extraAnnotations);
 
   // Trace re-collection from dump logic
   // Given a dump path, verifies it and re-collects the data into trace
@@ -65,6 +66,7 @@ class MmapBufferTraceWriter : public fbjni::HybridClass<MmapBufferTraceWriter> {
       const std::string& trace_prefix,
       int32_t trace_flags,
       std::shared_ptr<TraceCallbacks> callbacks,
+      const std::vector<std::pair<std::string, std::string>>& extraAnnotations,
       uint64_t timestamp = monotonicTime());
 
  private:

cpp/test/mmapbuf/writer/MmapBufferTraceWriterTest.cpp

@@ -13,6 +13,7 @@
 #include <fstream>
 #include <ostream>
 #include <sstream>
+#include <utility>
 #include <vector>
 
 #include <profilo/entries/Entry.h>
@@ -232,6 +233,8 @@ TEST_F(MmapBufferTraceWriterTest, testDumpWriteAndRecollectEndToEnd) {
   std::string testFolder(traceFolderPath());
   std::string testTracePrefix(kTracePrefix);
   auto mockCallbacks = std::make_shared<::testing::NiceMock<MockCallbacks>>();
+  std::vector<std::pair<std::string, std::string>> extraAnnotations{};
+
   MmapBufferTraceWriter traceWriter{};
 
   EXPECT_CALL(*mockCallbacks, onTraceStart(kTraceId, 0));
@@ -245,6 +248,7 @@ TEST_F(MmapBufferTraceWriterTest, testDumpWriteAndRecollectEndToEnd) {
       testTracePrefix,
       0,
       mockCallbacks,
+      extraAnnotations,
       kTraceRecollectionTimestamp);
 
   verifyLogEntriesFromTraceFile();
@@ -260,6 +264,8 @@ TEST_F(
   std::string testFolder(traceFolderPath());
   std::string testTracePrefix(kTracePrefix);
   auto mockCallbacks = std::make_shared<::testing::NiceMock<MockCallbacks>>();
+  std::vector<std::pair<std::string, std::string>> extraAnnotations{};
+
   MmapBufferTraceWriter traceWriter{};
 
   EXPECT_CALL(*mockCallbacks, onTraceStart(kTraceId, 0));
@@ -273,6 +279,7 @@ TEST_F(
       testTracePrefix,
       0,
       mockCallbacks,
+      extraAnnotations,
       kTraceRecollectionTimestamp);
 
   verifyLogEntriesFromTraceFile();
@@ -289,8 +296,9 @@ TEST_F(
   // exception.
   std::string testFolder("/dev/null");
   std::string testTracePrefix(kTracePrefix);
-
   auto mockCallbacks = std::make_shared<::testing::NiceMock<MockCallbacks>>();
+  std::vector<std::pair<std::string, std::string>> extraAnnotations{};
+
   MmapBufferTraceWriter traceWriter{};
 
   EXPECT_CALL(*mockCallbacks, onTraceAbort(kTraceId, AbortReason::UNKNOWN));
@@ -303,6 +311,7 @@ TEST_F(
       testTracePrefix,
       0,
       mockCallbacks,
+      extraAnnotations,
       kTraceRecollectionTimestamp);
 }
 
@@ -315,6 +324,8 @@ TEST_F(
   std::string testFolder(traceFolderPath());
   std::string testTracePrefix(kTracePrefix);
   auto mockCallbacks = std::make_shared<::testing::NiceMock<MockCallbacks>>();
+  std::vector<std::pair<std::string, std::string>> extraAnnotations{};
+
   MmapBufferTraceWriter traceWriter{};
 
   bool caughtException = false;
@@ -327,6 +338,7 @@ TEST_F(
         testTracePrefix,
         0,
         mockCallbacks,
+        extraAnnotations,
         kTraceRecollectionTimestamp);
   } catch (std::runtime_error& e) {
     ASSERT_STREQ(e.what(), "Unable to read the file-backed buffer.");

java/main/com/facebook/profilo/mmapbuf/writer/MmapBufferTraceWriter.java

@@ -44,5 +44,6 @@ public native void nativeWriteTrace(
       String traceFolder,
       String tracePrefix,
       int traceFlags,
-      NativeTraceWriterCallbacks mCallbacks);
+      NativeTraceWriterCallbacks mCallbacks,
+      String[] extraAnnotations);
 }