[trace][intel pt] Create a common accessor for live and postmortem data

Some parts of the code have to distinguish between live and postmortem threads
to figure out how to get some data, e.g. thread trace buffers. This makes the
code less generic and more error prone. An example of that is that we have
two different decoders: LiveThreadDecoder and PostMortemThreadDecoder. They
exist because getting the trace bufer is different for each case.

The problem doesn't stop there. Soon we'll have even more kinds of data, like
the context switch trace, whose fetching will be different for live and post-
mortem processes.

As a way to fix this, I'm creating a common API for accessing thread data,
which is able to figure out how to handle the postmortem and live cases on
behalf of the caller. As a result of that, I was able to eliminate the two
decoders and unify them into a simpler one. Not only that, our TraceSave
functionality only worked for live threads, but now it can also work for
postmortem processes, which might be useful now, but it might in the future.

This common API is OnThreadBinaryDataRead. More information in the inline
documentation.

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

lldb/include/lldb/Target/Trace.h

@@ -233,15 +233,76 @@ class Trace : public PluginInterface,
   ///     \a llvm::Error otherwise.
   llvm::Error Stop();
 
-  /// Get the trace file of the given post mortem thread.
-  llvm::Expected<const FileSpec &> GetPostMortemTraceFile(lldb::tid_t tid);
-
   /// \return
   ///     The stop ID of the live process being traced, or an invalid stop ID
   ///     if the trace is in an error or invalid state.
   uint32_t GetStopID();
 
+  using OnBinaryDataReadCallback =
+      std::function<llvm::Error(llvm::ArrayRef<uint8_t> data)>;
+  /// Fetch binary data associated with a thread, either live or postmortem, and
+  /// pass it to the given callback. The reason of having a callback is to free
+  /// the caller from having to manage the life cycle of the data and to hide
+  /// the different data fetching procedures that exist for live and post mortem
+  /// threads.
+  ///
+  /// The fetched data is not persisted after the callback is invoked.
+  ///
+  /// \param[in] tid
+  ///     The tid who owns the data.
+  ///
+  /// \param[in] kind
+  ///     The kind of data to read.
+  ///
+  /// \param[in] callback
+  ///     The callback to be invoked once the data was successfully read. Its
+  ///     return value, which is an \a llvm::Error, is returned by this
+  ///     function.
+  ///
+  /// \return
+  ///     An \a llvm::Error if the data couldn't be fetched, or the return value
+  ///     of the callback, otherwise.
+  llvm::Error OnThreadBinaryDataRead(lldb::tid_t tid, llvm::StringRef kind,
+                                     OnBinaryDataReadCallback callback);
+
 protected:
+  /// Implementation of \a OnThreadBinaryDataRead() for live threads.
+  llvm::Error OnLiveThreadBinaryDataRead(lldb::tid_t tid, llvm::StringRef kind,
+                                         OnBinaryDataReadCallback callback);
+
+  /// Implementation of \a OnThreadBinaryDataRead() for post mortem threads.
+  llvm::Error
+  OnPostMortemThreadBinaryDataRead(lldb::tid_t tid, llvm::StringRef kind,
+                                   OnBinaryDataReadCallback callback);
+
+  /// Get the file path containing data of a postmortem thread given a data
+  /// identifier.
+  ///
+  /// \param[in] tid
+  ///     The thread whose data is requested.
+  ///
+  /// \param[in] kind
+  ///     The kind of data requested.
+  ///
+  /// \return
+  ///     The file spec containing the requested data, or an \a llvm::Error in
+  ///     case of failures.
+  llvm::Expected<FileSpec> GetPostMortemThreadDataFile(lldb::tid_t tid,
+                                                       llvm::StringRef kind);
+
+  /// Associate a given thread with a data file using a data identifier.
+  ///
+  /// \param[in] tid
+  ///     The thread associated with the data file.
+  ///
+  /// \param[in] kind
+  ///     The kind of data being registered.
+  ///
+  /// \param[in] file_spec
+  ///     The path of the data file.
+  void SetPostMortemThreadDataFile(lldb::tid_t tid, llvm::StringRef kind,
+                                   FileSpec file_spec);
+
   /// Get binary data of a live thread given a data identifier.
   ///
   /// \param[in] tid
@@ -315,11 +376,25 @@ class Trace : public PluginInterface,
   uint32_t m_stop_id = LLDB_INVALID_STOP_ID;
   /// Process traced by this object if doing live tracing. Otherwise it's null.
   Process *m_live_process = nullptr;
+
+  /// These data kinds are returned by lldb-server when fetching the state of
+  /// the tracing session. The size in bytes can be used later for fetching the
+  /// data in batches.
+  /// \{
+
   /// tid -> data kind -> size
-  std::map<lldb::tid_t, std::unordered_map<std::string, size_t>>
+  llvm::DenseMap<lldb::tid_t, std::unordered_map<std::string, size_t>>
       m_live_thread_data;
+
   /// data kind -> size
   std::unordered_map<std::string, size_t> m_live_process_data;
+  /// \}
+
+  /// Postmortem traces can specific additional data files, which are
+  /// represented in this variable using a data kind identifier for each file.
+  /// tid -> data kind -> file
+  llvm::DenseMap<lldb::tid_t, std::unordered_map<std::string, FileSpec>>
+      m_postmortem_thread_data;
 };
 
 } // namespace lldb_private

lldb/source/Plugins/Trace/common/TraceSessionSaver.cpp

@@ -39,15 +39,12 @@ llvm::Error TraceSessionSaver::WriteSessionToFile(
 }
 
 llvm::Expected<JSONTraceSessionBase> TraceSessionSaver::BuildProcessesSection(
-    Process &live_process,
-    std::function<
-        llvm::Expected<llvm::Optional<std::vector<uint8_t>>>(lldb::tid_t tid)>
-        raw_trace_fetcher,
+    Process &live_process, llvm::StringRef raw_thread_trace_data_kind,
     FileSpec directory) {
 
   JSONTraceSessionBase json_session_description;
   Expected<std::vector<JSONThread>> json_threads =
-      BuildThreadsSection(live_process, raw_trace_fetcher, directory);
+      BuildThreadsSection(live_process, raw_thread_trace_data_kind, directory);
   if (!json_threads)
     return json_threads.takeError();
 
@@ -64,39 +61,41 @@ llvm::Expected<JSONTraceSessionBase> TraceSessionSaver::BuildProcessesSection(
 }
 
 llvm::Expected<std::vector<JSONThread>> TraceSessionSaver::BuildThreadsSection(
-    Process &live_process,
-    std::function<
-        llvm::Expected<llvm::Optional<std::vector<uint8_t>>>(lldb::tid_t tid)>
-        raw_trace_fetcher,
+    Process &live_process, llvm::StringRef raw_thread_trace_data_kind,
     FileSpec directory) {
   std::vector<JSONThread> json_threads;
   for (ThreadSP thread_sp : live_process.Threads()) {
+    TraceSP trace_sp = live_process.GetTarget().GetTrace();
+    lldb::tid_t tid = thread_sp->GetID();
+    if (!trace_sp->IsTraced(tid))
+      continue;
+
     // resolve the directory just in case
     FileSystem::Instance().Resolve(directory);
     FileSpec raw_trace_path = directory;
-    raw_trace_path.AppendPathComponent(std::to_string(thread_sp->GetID()) +
-                                       ".trace");
-    json_threads.push_back(JSONThread{static_cast<int64_t>(thread_sp->GetID()),
+    raw_trace_path.AppendPathComponent(std::to_string(tid) + ".trace");
+    json_threads.push_back(JSONThread{static_cast<int64_t>(tid),
                                       raw_trace_path.GetPath().c_str()});
 
-    llvm::Expected<llvm::Optional<std::vector<uint8_t>>> raw_trace =
-        raw_trace_fetcher(thread_sp->GetID());
-
-    if (!raw_trace)
-      return raw_trace.takeError();
-    if (!raw_trace.get())
-      continue;
-
-    std::basic_fstream<char> raw_trace_fs = std::fstream(
-        raw_trace_path.GetPath().c_str(), std::ios::out | std::ios::binary);
-    raw_trace_fs.write(reinterpret_cast<const char *>(&raw_trace.get()->at(0)),
-                       raw_trace.get()->size() * sizeof(uint8_t));
-    raw_trace_fs.close();
-    if (!raw_trace_fs) {
-      return createStringError(inconvertibleErrorCode(),
-                               formatv("couldn't write to the file {0}",
-                                       raw_trace_path.GetPath().c_str()));
-    }
+    llvm::Error err =
+        live_process.GetTarget().GetTrace()->OnThreadBinaryDataRead(
+            tid, raw_thread_trace_data_kind,
+            [&](llvm::ArrayRef<uint8_t> data) -> llvm::Error {
+              std::basic_fstream<char> raw_trace_fs =
+                  std::fstream(raw_trace_path.GetPath().c_str(),
+                               std::ios::out | std::ios::binary);
+              raw_trace_fs.write(reinterpret_cast<const char *>(&data[0]),
+                                 data.size() * sizeof(uint8_t));
+              raw_trace_fs.close();
+              if (!raw_trace_fs)
+                return createStringError(
+                    inconvertibleErrorCode(),
+                    formatv("couldn't write to the file {0}",
+                            raw_trace_path.GetPath().c_str()));
+              return Error::success();
+            });
+    if (err)
+      return std::move(err);
   }
   return json_threads;
 }

lldb/source/Plugins/Trace/common/TraceSessionSaver.h

@@ -39,24 +39,19 @@ class TraceSessionSaver {
   /// \param[in] live_process
   ///     The process being traced.
   ///
-  /// \param[in] raw_trace_fetcher
-  ///     Callback function that receives a thread ID and returns its raw trace.
-  ///     This callback should return \a None if the thread is not being traced.
-  ///     Otherwise, it should return the raw trace in bytes or an
-  ///     \a llvm::Error in case of failures.
+  /// \param[in] raw_thread_trace_data_kind
+  ///     Identifier for the data kind of the raw trace for each thread.
   ///
   /// \param[in] directory
   ///     The directory where files will be saved when building the processes
   ///     section.
   ///
   /// \return
   ///     The processes section or \a llvm::Error in case of failures.
-  static llvm::Expected<JSONTraceSessionBase> BuildProcessesSection(
-      Process &live_process,
-      std::function<
-          llvm::Expected<llvm::Optional<std::vector<uint8_t>>>(lldb::tid_t tid)>
-          raw_trace_fetcher,
-      FileSpec directory);
+  static llvm::Expected<JSONTraceSessionBase>
+  BuildProcessesSection(Process &live_process,
+                        llvm::StringRef raw_thread_trace_data_kind,
+                        FileSpec directory);
 
   /// Build the threads sub-section of the trace session description file.
   /// For each traced thread, its raw trace is also written to the file
@@ -65,24 +60,19 @@ class TraceSessionSaver {
   /// \param[in] live_process
   ///     The process being traced.
   ///
-  /// \param[in] raw_trace_fetcher
-  ///     Callback function that receives a thread ID and returns its raw trace.
-  ///     This callback should return \a None if the thread is not being traced.
-  ///     Otherwise, it should return the raw trace in bytes or an
-  ///     \a llvm::Error in case of failures.
+  /// \param[in] raw_thread_trace_data_kind
+  ///     Identifier for the data kind of the raw trace for each thread.
   ///
   /// \param[in] directory
   ///     The directory where files will be saved when building the threads
   ///     section.
   ///
   /// \return
   ///     The threads section or \a llvm::Error in case of failures.
-  static llvm::Expected<std::vector<JSONThread>> BuildThreadsSection(
-      Process &live_process,
-      std::function<
-          llvm::Expected<llvm::Optional<std::vector<uint8_t>>>(lldb::tid_t tid)>
-          raw_trace_fetcher,
-      FileSpec directory);
+  static llvm::Expected<std::vector<JSONThread>>
+  BuildThreadsSection(Process &live_process,
+                      llvm::StringRef raw_thread_trace_data_kind,
+                      FileSpec directory);
 
   /// Build modules sub-section of the trace's session. The original modules
   /// will be copied over to the \a <directory/modules> folder. Invalid modules

lldb/source/Plugins/Trace/intel-pt/LibiptDecoder.cpp

@@ -255,7 +255,7 @@ using PtInsnDecoderUP =
 static Expected<PtInsnDecoderUP>
 CreateInstructionDecoder(DecodedThread &decoded_thread,
                          TraceIntelPT &trace_intel_pt,
-                         MutableArrayRef<uint8_t> buffer) {
+                         ArrayRef<uint8_t> buffer) {
   Expected<pt_cpu> cpu_info = trace_intel_pt.GetCPUInfo();
   if (!cpu_info)
     return cpu_info.takeError();
@@ -268,8 +268,10 @@ CreateInstructionDecoder(DecodedThread &decoded_thread,
   if (IsLibiptError(status = pt_cpu_errata(&config.errata, &config.cpu)))
     return make_error<IntelPTError>(status);
 
-  config.begin = buffer.data();
-  config.end = buffer.data() + buffer.size();
+  // The libipt library does not modify the trace buffer, hence the
+  // following casts are safe.
+  config.begin = const_cast<uint8_t *>(buffer.data());
+  config.end = const_cast<uint8_t *>(buffer.data() + buffer.size());
 
   pt_insn_decoder *decoder_ptr = pt_insn_alloc_decoder(&config);
   if (!decoder_ptr)
@@ -285,9 +287,11 @@ CreateInstructionDecoder(DecodedThread &decoded_thread,
   return decoder_up;
 }
 
-void lldb_private::trace_intel_pt::DecodeTrace(
-    DecodedThread &decoded_thread, TraceIntelPT &trace_intel_pt,
-    MutableArrayRef<uint8_t> buffer) {
+void lldb_private::trace_intel_pt::DecodeTrace(DecodedThread &decoded_thread,
+                                               TraceIntelPT &trace_intel_pt,
+                                               ArrayRef<uint8_t> buffer) {
+  decoded_thread.SetRawTraceSize(buffer.size());
+
   Expected<PtInsnDecoderUP> decoder_up =
       CreateInstructionDecoder(decoded_thread, trace_intel_pt, buffer);
   if (!decoder_up)

lldb/source/Plugins/Trace/intel-pt/LibiptDecoder.h

@@ -21,7 +21,7 @@ namespace trace_intel_pt {
 /// instructions and errors in \p decoded_thread. It uses the low level libipt
 /// library underneath.
 void DecodeTrace(DecodedThread &decoded_thread, TraceIntelPT &trace_intel_pt,
-                 llvm::MutableArrayRef<uint8_t> buffer);
+                 llvm::ArrayRef<uint8_t> buffer);
 
 } // namespace trace_intel_pt
 } // namespace lldb_private

lldb/source/Plugins/Trace/intel-pt/ThreadDecoder.cpp

@@ -20,60 +20,25 @@ using namespace lldb_private;
 using namespace lldb_private::trace_intel_pt;
 using namespace llvm;
 
-// ThreadDecoder ====================
+ThreadDecoder::ThreadDecoder(const ThreadSP &thread_sp, TraceIntelPT &trace)
+    : m_thread_sp(thread_sp), m_trace(trace) {}
 
 DecodedThreadSP ThreadDecoder::Decode() {
   if (!m_decoded_thread.hasValue())
     m_decoded_thread = DoDecode();
   return *m_decoded_thread;
 }
 
-// LiveThreadDecoder ====================
-
-LiveThreadDecoder::LiveThreadDecoder(Thread &thread, TraceIntelPT &trace)
-    : m_thread_sp(thread.shared_from_this()), m_trace(trace) {}
-
-DecodedThreadSP LiveThreadDecoder::DoDecode() {
+DecodedThreadSP ThreadDecoder::DoDecode() {
   DecodedThreadSP decoded_thread_sp =
       std::make_shared<DecodedThread>(m_thread_sp);
 
-  Expected<std::vector<uint8_t>> buffer =
-      m_trace.GetLiveThreadBuffer(m_thread_sp->GetID());
-  if (!buffer) {
-    decoded_thread_sp->AppendError(buffer.takeError());
-    return decoded_thread_sp;
-  }
-
-  decoded_thread_sp->SetRawTraceSize(buffer->size());
-  DecodeTrace(*decoded_thread_sp, m_trace, MutableArrayRef<uint8_t>(*buffer));
-  return decoded_thread_sp;
-}
-
-// PostMortemThreadDecoder =======================
-
-PostMortemThreadDecoder::PostMortemThreadDecoder(
-    const lldb::ThreadPostMortemTraceSP &trace_thread, TraceIntelPT &trace)
-    : m_trace_thread(trace_thread), m_trace(trace) {}
-
-DecodedThreadSP PostMortemThreadDecoder::DoDecode() {
-  DecodedThreadSP decoded_thread_sp =
-      std::make_shared<DecodedThread>(m_trace_thread);
-
-  ErrorOr<std::unique_ptr<MemoryBuffer>> trace_or_error =
-      MemoryBuffer::getFile(m_trace_thread->GetTraceFile().GetPath());
-  if (std::error_code err = trace_or_error.getError()) {
-    decoded_thread_sp->AppendError(errorCodeToError(err));
-    return decoded_thread_sp;
-  }
-
-  MemoryBuffer &trace = **trace_or_error;
-  MutableArrayRef<uint8_t> trace_data(
-      // The libipt library does not modify the trace buffer, hence the
-      // following cast is safe.
-      reinterpret_cast<uint8_t *>(const_cast<char *>(trace.getBufferStart())),
-      trace.getBufferSize());
-  decoded_thread_sp->SetRawTraceSize(trace_data.size());
-
-  DecodeTrace(*decoded_thread_sp, m_trace, trace_data);
+  Error err = m_trace.OnThreadBufferRead(
+      m_thread_sp->GetID(), [&](llvm::ArrayRef<uint8_t> data) {
+        DecodeTrace(*decoded_thread_sp, m_trace, data);
+        return Error::success();
+      });
+  if (err)
+    decoded_thread_sp->AppendError(std::move(err));
   return decoded_thread_sp;
 }