[lldb] Cache Task pointer location and Task names

lldb/include/lldb/Target/Target.h

@@ -201,6 +201,8 @@ class TargetProperties : public Properties {
 
   bool GetSwiftUseTasksPlugin() const;
 
+  bool GetSwiftCacheTaskPtrLocation() const;
+
   Args GetSwiftPluginServerForPath() const;
 
   bool GetSwiftAutoImportFrameworks() const;

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,29 @@ 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::createStringError("could not get current task from thread: %s",
+                                   status.AsCString());
+  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 +2971,48 @@ 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
+    return ReadTaskAddr(it->second, *thread.GetProcess());
+  }
+
+  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 &&
+      real_thread.GetProcess()->GetTarget().GetSwiftCacheTaskPtrLocation())
+    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

@@ -38,6 +38,55 @@ void OperatingSystemSwiftTasks::Terminate() {
   PluginManager::UnregisterPlugin(CreateInstance);
 }
 
+/// A wrapper around ThreadMemory providing lazy name evaluation, as this is
+/// expensive to compute for Swift Tasks.
+class SwiftTaskThreadMemory : public ThreadMemory {
+public:
+  SwiftTaskThreadMemory(lldb_private::Process &process, lldb::tid_t tid,
+                        lldb::addr_t register_data_addr)
+      : ThreadMemory(process, tid, register_data_addr) {}
+
+  /// Updates the backing thread of this Task, as well as the location where the
+  /// task pointer is stored.
+  void UpdateBackingThread(const ThreadSP &new_backing_thread,
+                           lldb::addr_t task_addr) {
+    SetBackingThread(new_backing_thread);
+    m_task_addr = task_addr;
+  }
+
+  const char *GetName() override {
+    if (m_task_name.empty())
+      m_task_name = FindTaskName();
+    return m_task_name.c_str();
+  }
+
+private:
+  std::string GetDefaultTaskName() const {
+    return llvm::formatv("Task {0}", GetID());
+  }
+
+  /// If possible, read a user-provided task name from memory, otherwise use a
+  /// default name. This never returns an empty string.
+  std::string FindTaskName() const {
+    llvm::Expected<std::optional<std::string>> task_name =
+        GetTaskName(m_task_addr, *GetProcess());
+    if (auto err = task_name.takeError()) {
+      LLDB_LOG_ERROR(GetLog(LLDBLog::OS), std::move(err),
+                     "OperatingSystemSwiftTasks: failed while looking for name "
+                     "of task {1:x}: {0}",
+                     m_task_addr);
+      return GetDefaultTaskName();
+    }
+
+    if (!task_name->has_value())
+      return GetDefaultTaskName();
+    return llvm::formatv("{0} (Task {1})", *task_name, GetID());
+  }
+
+  std::string m_task_name = "";
+  lldb::addr_t m_task_addr = LLDB_INVALID_ADDRESS;
+};
+
 OperatingSystem *OperatingSystemSwiftTasks::CreateInstance(Process *process,
                                                            bool force) {
   if (!process || !process->GetTarget().GetSwiftUseTasksPlugin())
@@ -78,9 +127,9 @@ OperatingSystemSwiftTasks::OperatingSystemSwiftTasks(
     lldb_private::Process &process)
     : OperatingSystem(&process) {}
 
-ThreadSP OperatingSystemSwiftTasks::FindOrCreateSwiftThread(
-    ThreadList &old_thread_list, uint64_t task_id,
-    std::optional<std::string> task_name) {
+ThreadSP
+OperatingSystemSwiftTasks::FindOrCreateSwiftThread(ThreadList &old_thread_list,
+                                                   uint64_t task_id) {
   // Mask higher bits to avoid conflicts with core thread IDs.
   uint64_t masked_task_id = 0x0000000f00000000 | task_id;
 
@@ -89,19 +138,14 @@ ThreadSP OperatingSystemSwiftTasks::FindOrCreateSwiftThread(
       IsOperatingSystemPluginThread(old_thread))
     return old_thread;
 
-  std::string name;
-  if (task_name)
-    name = llvm::formatv("{0} (Task {1})", *task_name, task_id);
-  else
-    name = llvm::formatv("Task {0}", task_id);
-
-  return std::make_shared<ThreadMemoryProvidingName>(*m_process, masked_task_id,
-                                                     /*register_data_addr*/ 0,
-                                                     name);
+  return std::make_shared<SwiftTaskThreadMemory>(*m_process, masked_task_id,
+                                                 /*register_data_addr*/ 0);
 }
 
-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 "
@@ -130,27 +174,15 @@ static std::optional<uint64_t> FindTaskId(addr_t task_addr, Process &process) {
   return task_id;
 }
 
-static std::optional<std::string> FindTaskName(addr_t task_addr,
-                                               Process &process) {
-  auto task_name_or_err = GetTaskName(task_addr, process);
-  if (auto err = task_name_or_err.takeError()) {
-    LLDB_LOG_ERROR(GetLog(LLDBLog::OS), std::move(err),
-                   "OperatingSystemSwiftTasks: failed while looking for name "
-                   "of task {1:x}: {0}",
-                   task_addr);
-    return {};
-  }
-  return *task_name_or_err;
-}
-
 bool OperatingSystemSwiftTasks::UpdateThreadList(ThreadList &old_thread_list,
                                                  ThreadList &core_thread_list,
                                                  ThreadList &new_thread_list) {
   Log *log = GetLog(LLDBLog::OS);
   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) {
@@ -170,9 +202,9 @@ bool OperatingSystemSwiftTasks::UpdateThreadList(ThreadList &old_thread_list,
       continue;
     }
 
-    ThreadSP swift_thread = FindOrCreateSwiftThread(
-        old_thread_list, *task_id, FindTaskName(*task_addr, *m_process));
-    swift_thread->SetBackingThread(real_thread);
+    ThreadSP swift_thread = FindOrCreateSwiftThread(old_thread_list, *task_id);
+    static_cast<SwiftTaskThreadMemory &>(*swift_thread)
+        .UpdateBackingThread(real_thread, *task_addr);
     new_thread_list.AddThread(swift_thread);
     LLDB_LOGF(log,
               "OperatingSystemSwiftTasks: mapping thread IDs: %" PRIx64

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 {
@@ -49,8 +50,11 @@ class OperatingSystemSwiftTasks : public OperatingSystem {
   /// If a thread for task_id had been created in the last stop, return it.
   /// Otherwise, create a new MemoryThread for it.
   lldb::ThreadSP FindOrCreateSwiftThread(ThreadList &old_thread_list,
-                                         uint64_t task_id,
-                                         std::optional<std::string> task_name);
+                                         uint64_t task_id);
+
+  /// A cache for task addr locations, which are expensive to compute but
+  /// immutable.
+  TaskInspector m_task_inspector;
 };
 } // namespace lldb_private
 

lldb/source/Target/Target.cpp

@@ -4529,6 +4529,18 @@ bool TargetProperties::GetSwiftUseTasksPlugin() const {
   return true;
 }
 
+bool TargetProperties::GetSwiftCacheTaskPtrLocation() const {
+  const Property *exp_property =
+      m_collection_sp->GetPropertyAtIndex(ePropertyExperimental);
+  OptionValueProperties *exp_values =
+      exp_property->GetValue()->GetAsProperties();
+  if (exp_values)
+    return exp_values
+        ->GetPropertyAtIndexAs<bool>(ePropertySwiftCacheTaskPtrLocation)
+        .value_or(true);
+  return true;
+}
+
 Args TargetProperties::GetSwiftPluginServerForPath() const {
   const uint32_t idx = ePropertySwiftPluginServerForPath;
 

lldb/source/Target/TargetProperties.td

@@ -34,6 +34,9 @@ let Definition = "target_experimental" in {
   def SwiftUseTasksPlugin: Property<"swift-tasks-plugin-enabled", "Boolean">,
     DefaultTrue,
     Desc<"Enables the swift plugin converting tasks into threads">;
+  def SwiftCacheTaskPtrLocation: Property<"swift-cache-task-ptr-location", "Boolean">,
+    DefaultTrue,
+    Desc<"Enables caching of task pointers inside the swift tasks plugin">;
 }
 
 let Definition = "target" in {