pw_thread_threadx: Add basic thread snapshotting

Adds an optional library for capturing ThreadX thread state in
proto snapshots.

Change-Id: I47dbc5a57faf3d4411255ea83f04f17bf0a4dc46
Reviewed-on: https://pigweed-review.googlesource.com/c/pigweed/pigweed/+/51405
Commit-Queue: Rob Mohr <mohrr@google.com>
Pigweed-Auto-Submit: Rob Mohr <mohrr@google.com>
Reviewed-by: Ewout van Bekkum <ewout@google.com>
diff --git a/pw_thread_threadx/BUILD b/pw_thread_threadx/BUILD
index 924f13d..add260b 100644
--- a/pw_thread_threadx/BUILD
+++ b/pw_thread_threadx/BUILD
@@ -159,3 +159,39 @@
         "//pw_thread:yield_facade",
     ],
 )
+
+pw_cc_library(
+    name = "util",
+    hdrs = [
+        "public/pw_thread_threadx/util.h"
+    ],
+    srcs = [
+        "util.cc"
+    ],
+    deps = [
+        "//pw_function",
+        "//pw_status",
+    ],
+    # TODO(pwbug/317): This should depend on ThreadX but our third parties
+    # currently do not have Bazel support.
+)
+
+pw_cc_library(
+    name = "snapshot",
+    hdrs = [
+        "public/pw_thread_threadx/snapshot.h"
+    ],
+    srcs = [
+        "snapshot.cc"
+    ],
+    deps = [
+        ":util",
+        "//pw_thread:protos",
+        "//pw_bytes",
+        "//pw_function",
+        "//pw_protobuf",
+        "//pw_status",
+    ],
+    # TODO(pwbug/317): This should depend on ThreadX but our third parties
+    # currently do not have Bazel support.
+)
diff --git a/pw_thread_threadx/BUILD.gn b/pw_thread_threadx/BUILD.gn
index 79d475a..dd7d7b3 100644
--- a/pw_thread_threadx/BUILD.gn
+++ b/pw_thread_threadx/BUILD.gn
@@ -139,6 +139,33 @@
   deps = [ "$dir_pw_thread:yield.facade" ]
 }
 
+pw_source_set("util") {
+  public_configs = [ ":public_include_path" ]
+  public_deps = [
+    "$dir_pw_third_party/threadx",
+    dir_pw_function,
+    dir_pw_status,
+  ]
+  public = [ "public/pw_thread_threadx/util.h" ]
+  sources = [ "util.cc" ]
+}
+
+pw_source_set("snapshot") {
+  public_configs = [ ":public_include_path" ]
+  public_deps = [
+    "$dir_pw_third_party/threadx",
+    "$dir_pw_thread:protos.pwpb",
+    "$dir_pw_thread:snapshot",
+    dir_pw_bytes,
+    dir_pw_function,
+    dir_pw_protobuf,
+    dir_pw_status,
+  ]
+  public = [ "public/pw_thread_threadx/snapshot.h" ]
+  sources = [ "snapshot.cc" ]
+  deps = [ ":util" ]
+}
+
 pw_test_group("tests") {
   tests = [ ":thread_backend_test" ]
 }
diff --git a/pw_thread_threadx/docs.rst b/pw_thread_threadx/docs.rst
index 527b487..5cfd84f 100644
--- a/pw_thread_threadx/docs.rst
+++ b/pw_thread_threadx/docs.rst
@@ -25,3 +25,43 @@
   * - ``pw_thread:thread``
     - ``pw_thread_threadx:thread``
     - Thread creation.
+
+---------
+utilities
+---------
+In cases where an operation must be performed for every thread,
+``ForEachThread()`` can be used to iterate over all the created thread TCBs.
+Note that it's only safe to use this while the scheduler is disabled.
+
+--------------------
+Snapshot integration
+--------------------
+This ``pw_thread`` backend provides helper functions that capture ThreadX thread
+state to a ``pw::thread::Thread`` proto.
+
+SnapshotThread()/SnapshotThreads()
+==================================
+``SnapshotThread()`` captures the thread name, state, and stack information for
+the provided RTX TCB to a ``pw::thread::Thread`` protobuf encoder. To ensure
+the most up-to-date information is captured, the stack pointer for the currently
+running thread must be provided for cases where the running thread is being
+captured. For ARM Cortex-M CPUs, you can do something like this:
+
+.. Code:: cpp
+
+  // Capture PSP.
+  void* stack_ptr = 0;
+  asm volatile("mrs %0, psp\n" : "=r"(stack_ptr));
+  pw::thread::ProcessThreadStackCallback cb =
+      [](pw::thread::Thread::StreamEncoder& encoder,
+         pw::ConstByteSpan stack) -> pw::Status {
+    return encoder.WriteRawStack(stack);
+  };
+  pw::thread::threadx::SnapshotThread(my_thread, stack_ptr,
+                                      snapshot_encoder, cb);
+
+``SnapshotThreads()`` wraps the singular thread capture to instead captures
+all created threads to a ``pw::thread::SnapshotThreadInfo`` message. This proto
+message overlays a snapshot, so it is safe to static cast a
+``pw::snapshot::Snapshot::StreamEncoder`` to a
+``pw::thread::SnapshotThreadInfo::StreamEncoder`` when calling this function.
diff --git a/pw_thread_threadx/public/pw_thread_threadx/snapshot.h b/pw_thread_threadx/public/pw_thread_threadx/snapshot.h
new file mode 100644
index 0000000..1845b53
--- /dev/null
+++ b/pw_thread_threadx/public/pw_thread_threadx/snapshot.h
@@ -0,0 +1,74 @@
+// Copyright 2021 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#pragma once
+
+#include "pw_protobuf/encoder.h"
+#include "pw_status/status.h"
+#include "pw_thread/snapshot.h"
+#include "pw_thread_protos/thread.pwpb.h"
+#include "tx_api.h"
+
+namespace pw::thread::threadx {
+
+// Captures all threadx threads in a system as part of a snapshot.
+//
+// An updated running_thread_stack_pointer must be provided in order for the
+// running thread's context to reflect the running state. For ARM, you might do
+// something like this:
+//
+//    // Capture PSP.
+//    void* stack_ptr = 0;
+//    asm volatile("mrs %0, psp\n" : "=r"(stack_ptr));
+//    pw::thread::ProcessThreadStackCallback cb =
+//        [](pw::thread::Thread::StreamEncoder& encoder,
+//           pw::ConstByteSpan stack) -> pw::Status {
+//      return encoder.WriteRawStack(stack);
+//    };
+//    pw::thread::threadx::SnapshotThread(my_thread, stack_ptr,
+//                                        snapshot_encoder, cb);
+//
+// Warning: This is only safe to use when interrupts and the scheduler are
+// disabled!
+// Warning: SMP ports are not yet supported.
+Status SnapshotThreads(void* running_thread_stack_pointer,
+                       SnapshotThreadInfo::StreamEncoder& encoder,
+                       ProcessThreadStackCallback& thread_stack_callback);
+
+// Captures only the provided thread handle as a pw::thread::Thread proto
+// message. After thread info capture, the ProcessThreadStackCallback is called
+// to capture either the raw_stack or raw_backtrace.
+//
+// An updated running_thread_stack_pointer must be provided in order for the
+// running thread's context to reflect the current state. If the thread being
+// captured is not the running thread, the value is ignored. Note that the
+// stack pointer in the thread handle is almost always stale on the running
+// thread.
+//
+// Captures the following proto fields:
+//   pw.thread.Thread:
+//     name
+//     state
+//     stack_start_pointer
+//     stack_end_pointer
+//     stack_pointer
+//
+// Warning: This is only safe to use when interrupts and the scheduler are
+// disabled!
+// Warning: SMP ports are not yet supported.
+Status SnapshotThread(const TX_THREAD& thread,
+                      void* running_thread_stack_pointer,
+                      Thread::StreamEncoder& encoder,
+                      ProcessThreadStackCallback& thread_stack_callback);
+
+}  // namespace pw::thread::threadx
diff --git a/pw_thread_threadx/public/pw_thread_threadx/util.h b/pw_thread_threadx/public/pw_thread_threadx/util.h
new file mode 100644
index 0000000..ed0c1e1
--- /dev/null
+++ b/pw_thread_threadx/public/pw_thread_threadx/util.h
@@ -0,0 +1,39 @@
+// Copyright 2021 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#pragma once
+
+#include "pw_function/function.h"
+#include "pw_status/status.h"
+#include "tx_api.h"
+
+namespace pw::thread::threadx {
+
+// A callback that is executed for each thread when using ForEachThread().
+using ThreadCallback = pw::Function<Status(const TX_THREAD&)>;
+
+// Iterates through all threads that haven't been deleted, calling the provided
+// callback on each thread. If the callback fails on one thread, the iteration
+// stops.
+//
+// Warning: This is only safe to use when the scheduler is disabled.
+Status ForEachThread(ThreadCallback& cb);
+
+namespace internal {
+
+// This function is exposed for testing. Prefer
+// pw::thread::threadx::ForEachThread.
+Status ForEachThread(const TX_THREAD& starting_thread, ThreadCallback& cb);
+
+}  // namespace internal
+}  // namespace pw::thread::threadx
diff --git a/pw_thread_threadx/snapshot.cc b/pw_thread_threadx/snapshot.cc
new file mode 100644
index 0000000..0333714
--- /dev/null
+++ b/pw_thread_threadx/snapshot.cc
@@ -0,0 +1,131 @@
+// Copyright 2021 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#include "pw_thread_threadx/snapshot.h"
+
+#include <string_view>
+
+#include "pw_function/function.h"
+#include "pw_protobuf/encoder.h"
+#include "pw_status/status.h"
+#include "pw_thread/snapshot.h"
+#include "pw_thread_protos/thread.pwpb.h"
+#include "pw_thread_threadx/util.h"
+#include "tx_api.h"
+#include "tx_thread.h"
+
+namespace pw::thread::threadx {
+namespace {
+
+// TODO(amontanez): This might make unit testing codepaths that use this more
+// challenging.
+inline bool ThreadIsRunning(const TX_THREAD& thread) {
+  const TX_THREAD* running_thread;
+  TX_THREAD_GET_CURRENT(running_thread);
+  return running_thread == &thread;
+}
+
+void CaptureThreadState(const TX_THREAD& thread,
+                        Thread::StreamEncoder& encoder) {
+  if (ThreadIsRunning(thread)) {
+    encoder.WriteState(ThreadState::Enum::RUNNING);
+    return;
+  }
+
+  switch (thread.tx_thread_state) {
+    case TX_READY:
+      encoder.WriteState(ThreadState::Enum::READY);
+      break;
+    case TX_COMPLETED:
+    case TX_TERMINATED:
+      encoder.WriteState(ThreadState::Enum::INACTIVE);
+      break;
+    case TX_SUSPENDED:
+    case TX_SLEEP:
+      encoder.WriteState(ThreadState::Enum::SUSPENDED);
+      break;
+    case TX_QUEUE_SUSP:
+    case TX_SEMAPHORE_SUSP:
+    case TX_EVENT_FLAG:
+    case TX_BLOCK_MEMORY:
+    case TX_BYTE_MEMORY:
+    case TX_IO_DRIVER:
+    case TX_FILE:
+    case TX_TCP_IP:
+    case TX_MUTEX_SUSP:
+      encoder.WriteState(ThreadState::Enum::BLOCKED);
+      break;
+    default:
+      encoder.WriteState(ThreadState::Enum::UNKNOWN);
+  }
+}
+
+}  // namespace
+
+Status SnapshotThreads(void* running_thread_stack_pointer,
+                       SnapshotThreadInfo::StreamEncoder& encoder,
+                       ProcessThreadStackCallback& stack_dumper) {
+  struct {
+    void* running_thread_stack_pointer;
+    SnapshotThreadInfo::StreamEncoder* encoder;
+    ProcessThreadStackCallback* stack_dumper;
+  } ctx;
+  ctx.running_thread_stack_pointer = running_thread_stack_pointer;
+  ctx.encoder = &encoder;
+  ctx.stack_dumper = &stack_dumper;
+
+  ThreadCallback thread_capture_cb([&ctx](const TX_THREAD& thread) -> Status {
+    Thread::StreamEncoder thread_encoder = ctx.encoder->GetThreadsEncoder();
+    return SnapshotThread(thread,
+                          ctx.running_thread_stack_pointer,
+                          thread_encoder,
+                          *ctx.stack_dumper);
+  });
+
+  return ForEachThread(thread_capture_cb);
+}
+
+Status SnapshotThread(const TX_THREAD& thread,
+                      void* running_thread_stack_pointer,
+                      Thread::StreamEncoder& encoder,
+                      ProcessThreadStackCallback& thread_stack_callback) {
+  encoder.WriteName(
+      std::as_bytes(std::span(std::string_view(thread.tx_thread_name))));
+
+  CaptureThreadState(thread, encoder);
+
+  const StackContext thread_ctx = {
+      .thread_name = thread.tx_thread_name,
+
+      // TODO(amontanez): When ThreadX is built with stack checking enabled, the
+      // lowest-addressed `unsigned long` is reserved for a watermark. This
+      // means in practice the stack pointer should never end up there. To be
+      // conservative, behave as though TX_THREAD_STACK_CHECK is always fully
+      // enabled.
+      .stack_low_addr =
+          reinterpret_cast<uintptr_t>(thread.tx_thread_stack_start) +
+          sizeof(ULONG),
+
+      .stack_high_addr =
+          reinterpret_cast<uintptr_t>(thread.tx_thread_stack_end),
+
+      // If the thread is active, the stack pointer in the TCB is stale.
+      .stack_pointer = reinterpret_cast<uintptr_t>(
+          ThreadIsRunning(thread) ? running_thread_stack_pointer
+                                  : thread.tx_thread_stack_ptr),
+  };
+
+  return SnapshotStack(thread_ctx, encoder, thread_stack_callback);
+}
+
+}  // namespace pw::thread::threadx
diff --git a/pw_thread_threadx/util.cc b/pw_thread_threadx/util.cc
new file mode 100644
index 0000000..0d9c5e5
--- /dev/null
+++ b/pw_thread_threadx/util.cc
@@ -0,0 +1,44 @@
+// Copyright 2021 The Pigweed Authors
+//
+// Licensed under the Apache License, Version 2.0 (the "License"); you may not
+// use this file except in compliance with the License. You may obtain a copy of
+// the License at
+//
+//     https://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+// WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+// License for the specific language governing permissions and limitations under
+// the License.
+#include "pw_thread_threadx/util.h"
+
+#include "pw_function/function.h"
+#include "pw_status/status.h"
+#include "pw_status/try.h"
+#include "tx_api.h"
+#include "tx_thread.h"
+
+namespace pw::thread::threadx {
+
+namespace internal {
+
+// Iterates through all threads that haven't been deleted, calling the provided
+// call
+Status ForEachThread(const TX_THREAD& starting_thread, ThreadCallback& cb) {
+  const TX_THREAD* thread = &starting_thread;
+  do {
+    PW_TRY(cb(*thread));
+    thread = thread->tx_thread_created_next;
+  } while (thread != &starting_thread);
+
+  return OkStatus();
+}
+
+}  // namespace internal
+
+Status ForEachThread(ThreadCallback& cb) {
+  return internal::ForEachThread(*_tx_thread_created_ptr, cb);
+}
+
+}  // namespace pw::thread::threadx