Add an API to iterate through Status payloads

This is useful when we need to pass status through API boundary.

PiperOrigin-RevId: 549475209

tensorflow/c/tf_status.cc

@@ -28,6 +28,10 @@ void TF_SetStatus(TF_Status* s, TF_Code code, const char* msg) {
 void TF_SetPayload(TF_Status* s, const char* key, const char* value) {
   TSL_SetPayload(s, key, value);
 }
+void TF_ForEachPayload(const TF_Status* s, TF_PayloadVisitor visitor,
+                       void* capture) {
+  TSL_ForEachPayload(s, visitor, capture);
+}
 void TF_SetStatusFromIOError(TF_Status* s, int error_code,
                              const char* context) {
   TSL_SetStatusFromIOError(s, error_code, context);

tensorflow/c/tf_status.h

@@ -66,6 +66,14 @@ TF_CAPI_EXPORT extern void TF_SetStatus(TF_Status* s, TF_Code code,
 TF_CAPI_EXPORT void TF_SetPayload(TF_Status* s, const char* key,
                                   const char* value);
 
+// Iterates over the stored payloads and calls the `visitor(key, value)`
+// callable for each one. `key` and `value` is only usable during the callback.
+// `capture` will be passed to the callback without modification.
+#define TF_PayloadVisitor TSL_PayloadVisitor
+TF_CAPI_EXPORT extern void TF_ForEachPayload(const TF_Status* s,
+                                             TF_PayloadVisitor visitor,
+                                             void* capture);
+
 // Convert from an I/O error code (e.g., errno) to a TF_Status value.
 // Any previous information is lost. Prefer to use this instead of TF_SetStatus
 // when the error comes from I/O operations.

tensorflow/tsl/c/tsl_status.cc

@@ -15,6 +15,8 @@ limitations under the License.
 
 #include "tensorflow/tsl/c/tsl_status.h"
 
+#include <string>
+
 #include "tensorflow/tsl/c/tsl_status_internal.h"
 #include "tensorflow/tsl/platform/errors.h"
 #include "tensorflow/tsl/platform/status.h"
@@ -40,6 +42,16 @@ void TSL_SetPayload(TSL_Status* s, const char* key, const char* value) {
   s->status.SetPayload(key, absl::Cord(absl::string_view(value)));
 }
 
+void TSL_ForEachPayload(const TSL_Status* s, TSL_PayloadVisitor visitor,
+                        void* capture) {
+  s->status.ForEachPayload([visitor, capture](absl::string_view type_url,
+                                              const absl::Cord& payload) {
+    std::string type_url_str(type_url);
+    std::string payload_str(payload);
+    visitor(type_url_str.c_str(), payload_str.c_str(), capture);
+  });
+}
+
 void TSL_SetStatusFromIOError(TSL_Status* s, int error_code,
                               const char* context) {
   // TODO(b/139060984): Handle windows when changing its filesystem

tensorflow/tsl/c/tsl_status.h

@@ -60,7 +60,15 @@ extern void TSL_SetStatus(TSL_Status* s, TSL_Code code, const char* msg);
 // Record <key, value> as a payload in *s. The previous payload having the
 // same key (if any) is overwritten. Payload will not be added if the Status
 // is OK.
-void TSL_SetPayload(TSL_Status* s, const char* key, const char* value);
+extern void TSL_SetPayload(TSL_Status* s, const char* key, const char* value);
+
+// Iterates over the stored payloads and calls the `visitor(key, value)`
+// callable for each one. `key` and `value` is only usable during the callback.
+// `capture` will be passed to the callback without modification.
+typedef void (*TSL_PayloadVisitor)(const char* key, const char* value,
+                                   void* capture);
+extern void TSL_ForEachPayload(const TSL_Status* s, TSL_PayloadVisitor visitor,
+                               void* capture);
 
 // Convert from an I/O error code (e.g., errno) to a TSL_Status value.
 // Any previous information is lost. Prefer to use this instead of TSL_SetStatus

tensorflow/tsl/c/tsl_status_test.cc

@@ -33,8 +33,15 @@ TEST(TSL_Status, PayloadsSet) {
   TSL_SetPayload(tsl_status, "b", "2");
   TSL_SetPayload(tsl_status, "c", "3");
 
-  const std::unordered_map<std::string, std::string> payloads =
-      errors::GetPayloads(tsl_status->status);
+  std::unordered_map<std::string, std::string> payloads;
+  TSL_ForEachPayload(
+      tsl_status,
+      [](const char* key, const char* value, void* capture) {
+        std::unordered_map<std::string, std::string>* payloads =
+            static_cast<std::unordered_map<std::string, std::string>*>(capture);
+        payloads->emplace(key, value);
+      },
+      &payloads);
   EXPECT_EQ(payloads.size(), 3);
   EXPECT_EQ(payloads.at("a"), "1");
   EXPECT_EQ(payloads.at("b"), "2");