[lldb][swift] Cache location of Task Pointers for OperatingSystemSwiftTasks

Obtaining the TLS from debugserver over a wired (or slower) conection
was identified as a major source of slowdown after introducing
OperatingSystemSwiftTasks. This is due to multiple memory reads being
necessary:

1. Memory reads to obtain auxiliary information from libpthread that is
   then sent to the jThreadExtendedInfo packet. See
   "plo_pthread_tsd_base_address_offset","plo_pthread_tsd_base_offset"
   and "plo_pthread_tsd_entry_size".
2. A packet requesting jThreadExtendedInfo.

Both 1 and 2 are done per thread, on every stop. However, there is
strong evidence that we can safely cache the TLS pointer after a certain
point in the thread lifetime, which is what this patch accomplishes.

Caching is done at the OperatingSystemSwiftTasks level, not at the
core Thread class. This is done for two reasons:

1. To minimize the surface area / risk of this change.
2. jThreadExtendedInfo can produce an invalid TLS address earlier in the
   thread lifetime, so we can only cache this if a memory read at that
   location is successful.

On local testing over a wired connection, this reduces stepping time by
about 100ms, from an average of 300ms.

Author: felipepiovezan

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.cpp

@@ -2325,8 +2325,9 @@ class CommandObjectLanguageSwiftTaskInfo final : public CommandObjectParsed {
         return;
       }
 
-      auto task_addr_or_err =
-          GetTaskAddrFromThreadLocalStorage(m_exe_ctx.GetThreadRef());
+      TaskInspector task_inspector;
+      auto task_addr_or_err = task_inspector.GetTaskAddrFromThreadLocalStorage(
+          m_exe_ctx.GetThreadRef());
       if (auto error = task_addr_or_err.takeError()) {
         result.AppendError(toString(std::move(error)));
         return;
@@ -2939,17 +2940,30 @@ std::optional<lldb::addr_t> SwiftLanguageRuntime::TrySkipVirtualParentProlog(
   return pc_value + prologue_size;
 }
 
-llvm::Expected<lldb::addr_t> GetTaskAddrFromThreadLocalStorage(Thread &thread) {
+/// Attempts to read the memory location at `task_addr_location`, producing
+/// the Task pointer if possible.
+static llvm::Expected<lldb::addr_t>
+ReadTaskAddr(lldb::addr_t task_addr_location, Process &process) {
+  Status status;
+  addr_t task_addr = process.ReadPointerFromMemory(task_addr_location, status);
+  if (status.Fail())
+    return llvm::joinErrors(
+        llvm::createStringError("could not get current task from thread"),
+        status.takeError());
+  return task_addr;
+}
+
+/// Compute the location where the Task pointer for `real_thread` is stored by
+/// the runtime.
+static llvm::Expected<lldb::addr_t>
+ComputeTaskAddrLocationFromThreadLocalStorage(Thread &real_thread) {
 #if !SWIFT_THREADING_USE_RESERVED_TLS_KEYS
   return llvm::createStringError(
       "getting the current task from a thread is not supported");
 #else
   // Compute the thread local storage address for this thread.
   addr_t tsd_addr = LLDB_INVALID_ADDRESS;
 
-  // Look through backing threads when inspecting TLS.
-  Thread &real_thread =
-      thread.GetBackingThread() ? *thread.GetBackingThread() : thread;
   if (auto info_sp = real_thread.GetExtendedInfo())
     if (auto *info_dict = info_sp->GetAsDictionary())
       info_dict->GetValueForKeyAsInteger("tsd_address", tsd_addr);
@@ -2958,18 +2972,55 @@ llvm::Expected<lldb::addr_t> GetTaskAddrFromThreadLocalStorage(Thread &thread) {
     return llvm::createStringError("could not read current task from thread");
 
   // Offset of the Task pointer in a Thread's local storage.
-  Process &process = *thread.GetProcess();
+  Process &process = *real_thread.GetProcess();
   size_t ptr_size = process.GetAddressByteSize();
   uint64_t task_ptr_offset_in_tls =
       swift::tls_get_key(swift::tls_key::concurrency_task) * ptr_size;
-  addr_t task_addr_location = tsd_addr + task_ptr_offset_in_tls;
-  Status status;
-  addr_t task_addr = process.ReadPointerFromMemory(task_addr_location, status);
-  if (status.Fail())
-    return llvm::createStringError("could not get current task from thread: %s",
-                                   status.AsCString());
-  return task_addr;
+  return tsd_addr + task_ptr_offset_in_tls;
+#endif
+}
+
+llvm::Expected<lldb::addr_t>
+TaskInspector::GetTaskAddrFromThreadLocalStorage(Thread &thread) {
+  // Look through backing threads when inspecting TLS.
+  Thread &real_thread =
+      thread.GetBackingThread() ? *thread.GetBackingThread() : thread;
+
+  if (auto it = m_tid_to_task_addr_location.find(real_thread.GetID());
+      it != m_tid_to_task_addr_location.end()) {
+#ifndef NDEBUG
+    // In assert builds, check that caching did not produce incorrect results.
+    llvm::Expected<lldb::addr_t> task_addr_location =
+        ComputeTaskAddrLocationFromThreadLocalStorage(real_thread);
+    assert(task_addr_location);
+    assert(it->second == *task_addr_location);
 #endif
+    llvm::Expected<lldb::addr_t> task_addr =
+        ReadTaskAddr(it->second, *thread.GetProcess());
+    if (task_addr)
+      return task_addr;
+    // If the cached task addr location became invalid, invalidate the cache.
+    m_tid_to_task_addr_location.erase(it);
+    LLDB_LOG_ERROR(GetLog(LLDBLog::OS), task_addr.takeError(),
+                   "TaskInspector: evicted task location address due to "
+                   "invalid memory read: {0}");
+  }
+
+  llvm::Expected<lldb::addr_t> task_addr_location =
+      ComputeTaskAddrLocationFromThreadLocalStorage(real_thread);
+  if (!task_addr_location)
+    return task_addr_location;
+
+  llvm::Expected<lldb::addr_t> task_addr =
+      ReadTaskAddr(*task_addr_location, *thread.GetProcess());
+
+  // If the read from this TLS address is successful, cache the TLS address.
+  // Caching without a valid read is dangerous: earlier in the thread
+  // lifetime, the result of GetExtendedInfo can be invalid.
+  if (task_addr)
+    m_tid_to_task_addr_location.try_emplace(real_thread.GetID(),
+                                            *task_addr_location);
+  return task_addr;
 }
 
 namespace {

lldb/source/Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.h

@@ -902,9 +902,17 @@ struct AsyncUnwindRegisterNumbers {
 std::optional<AsyncUnwindRegisterNumbers>
 GetAsyncUnwindRegisterNumbers(llvm::Triple::ArchType triple);
 
-/// Inspects thread local storage to find the address of the currently executing
-/// task.
-llvm::Expected<lldb::addr_t> GetTaskAddrFromThreadLocalStorage(Thread &thread);
+/// A helper class to find and cache the location of Task pointer inside TLS.
+class TaskInspector {
+public:
+  /// Inspects thread local storage to find the address of the currently
+  /// executing task, if any.
+  llvm::Expected<lldb::addr_t>
+  GetTaskAddrFromThreadLocalStorage(Thread &thread);
+
+private:
+  llvm::DenseMap<uint64_t, lldb::addr_t> m_tid_to_task_addr_location;
+};
 
 llvm::Expected<std::optional<std::string>> GetTaskName(lldb::addr_t task,
                                                        Process &process);

lldb/source/Plugins/OperatingSystem/SwiftTasks/OperatingSystemSwiftTasks.cpp

@@ -100,8 +100,10 @@ ThreadSP OperatingSystemSwiftTasks::FindOrCreateSwiftThread(
                                                      name);
 }
 
-static std::optional<addr_t> FindTaskAddress(Thread &thread) {
-  llvm::Expected<addr_t> task_addr = GetTaskAddrFromThreadLocalStorage(thread);
+static std::optional<addr_t> FindTaskAddress(TaskInspector &task_inspector,
+                                             Thread &thread) {
+  llvm::Expected<addr_t> task_addr =
+      task_inspector.GetTaskAddrFromThreadLocalStorage(thread);
   if (!task_addr) {
     LLDB_LOG_ERROR(GetLog(LLDBLog::OS), task_addr.takeError(),
                    "OperatingSystemSwiftTasks: failed to find task address in "
@@ -150,7 +152,8 @@ bool OperatingSystemSwiftTasks::UpdateThreadList(ThreadList &old_thread_list,
   LLDB_LOG(log, "OperatingSystemSwiftTasks: Updating thread list");
 
   for (const ThreadSP &real_thread : core_thread_list.Threads()) {
-    std::optional<addr_t> task_addr = FindTaskAddress(*real_thread);
+    std::optional<addr_t> task_addr =
+        FindTaskAddress(m_task_inspector, *real_thread);
 
     // If this is not a thread running a Task, add it to the list as is.
     if (!task_addr) {

lldb/source/Plugins/OperatingSystem/SwiftTasks/OperatingSystemSwiftTasks.h

@@ -11,6 +11,7 @@
 
 #if LLDB_ENABLE_SWIFT
 
+#include "Plugins/LanguageRuntime/Swift/SwiftLanguageRuntime.h"
 #include "lldb/Target/OperatingSystem.h"
 
 namespace lldb_private {
@@ -51,6 +52,10 @@ class OperatingSystemSwiftTasks : public OperatingSystem {
   lldb::ThreadSP FindOrCreateSwiftThread(ThreadList &old_thread_list,
                                          uint64_t task_id,
                                          std::optional<std::string> task_name);
+
+  /// A cache for task addr locations, which are expensive to compute but
+  /// immutable.
+  TaskInspector m_task_inspector;
 };
 } // namespace lldb_private