feat(status): support rocksdb::Status fowarding in GET_OR_RET (#2630)

apache · Oct 31, 2024 · a7d5f4f · a7d5f4f
1 parent 3a51869
commit a7d5f4f
Show file tree

Hide file tree

Showing 2 changed files with 55 additions and 6 deletions.
diff --git a/src/common/status.h b/src/common/status.h
@@ -29,6 +29,7 @@
 #include <type_traits>
 #include <utility>
 
+#include "rocksdb/status.h"
 #include "type_util.h"
 
 class [[nodiscard]] Status {
@@ -369,10 +370,30 @@ struct [[nodiscard]] StatusOr {
   friend struct StatusOr;
 };
 
+template <typename T,
+          std::enable_if_t<IsStatusOr<RemoveCVRef<T>>::value || std::is_same_v<RemoveCVRef<T>, Status>, int> = 0>
+decltype(auto) StatusGetValue(T&& v) {
+  return std::forward<T>(v).GetValue();
+}
+
+template <typename T, std::enable_if_t<std::is_same_v<RemoveCVRef<T>, rocksdb::Status>, int> = 0>
+void StatusGetValue(T&&) {}
+
+template <typename T,
+          std::enable_if_t<IsStatusOr<RemoveCVRef<T>>::value || std::is_same_v<RemoveCVRef<T>, Status>, int> = 0>
+bool StatusIsOK(const T& v) {
+  return v.IsOK();
+}
+
+template <typename T, std::enable_if_t<std::is_same_v<RemoveCVRef<T>, rocksdb::Status>, int> = 0>
+bool StatusIsOK(const T& v) {
+  return v.ok();
+}
+
 // NOLINTNEXTLINE
-#define GET_OR_RET(...)                                         \
-  ({                                                            \
-    auto&& status = (__VA_ARGS__);                              \
-    if (!status) return std::forward<decltype(status)>(status); \
-    std::forward<decltype(status)>(status);                     \
-  }).GetValue()
+#define GET_OR_RET(...)                                                     \
+  StatusGetValue(({                                                         \
+    auto&& status = (__VA_ARGS__);                                          \
+    if (!StatusIsOK(status)) return std::forward<decltype(status)>(status); \
+    std::forward<decltype(status)>(status);                                 \
+  }))
diff --git a/tests/cppunit/status_test.cc b/tests/cppunit/status_test.cc
@@ -226,3 +226,31 @@ TEST(StatusOr, Prefixed) {
   ASSERT_EQ(g(-2).Msg(), "oh: hi");
   ASSERT_EQ(*g(5), 36);
 }
+
+TEST(GetOrRet, RocksdbStatus) {
+  auto f = [](int x) -> Status {
+    if (x < 10) return {Status::NotOK};
+    return Status::OK();
+  };
+
+  auto g = [&f](int x) -> Status {
+    GET_OR_RET(f(x));
+    return Status::OK();
+  };
+
+  ASSERT_TRUE(g(10));
+  ASSERT_FALSE(g(1));
+
+  auto f2 = [](int x) -> rocksdb::Status {
+    if (x < 10) return rocksdb::Status::InvalidArgument("");
+    return rocksdb::Status::OK();
+  };
+
+  auto g2 = [&f2](int x) -> rocksdb::Status {
+    GET_OR_RET(f2(x));
+    return rocksdb::Status::OK();
+  };
+
+  ASSERT_TRUE(g2(10).ok());
+  ASSERT_FALSE(g2(1).ok());
+}