add macro VELOX_RETURN_UNEXPECTED*

marin-ma · Nov 18, 2024 · 28f66a9 · 28f66a9
1 parent 3eb2fa4
commit 28f66a9
Show file tree

Hide file tree

Showing 2 changed files with 65 additions and 1 deletion.
diff --git a/velox/common/base/Status.h b/velox/common/base/Status.h
@@ -485,6 +485,29 @@ void Status::moveFrom(Status& s) {
     VELOX_RETURN_IF(!__s.ok(), __s);                          \
   } while (false)
 
+/// Return with given status wrapped in folly::Unexpected if condition is met.
+#define VELOX_RETURN_UNEXPECTED_IF(condition, status) \
+  do {                                                \
+    if (FOLLY_UNLIKELY(condition)) {                  \
+      return (::folly::makeUnexpected(status));       \
+    }                                                 \
+  } while (false)
+
+/// Propagate any non-successful Status wrapped in folly::Unexpected to the
+/// caller.
+#define VELOX_RETURN_UNEXPECTED_NOT_OK(status)                \
+  do {                                                        \
+    ::facebook::velox::Status __s =                           \
+        ::facebook::velox::internal::genericToStatus(status); \
+    VELOX_RETURN_IF(!__s.ok(), ::folly::makeUnexpected(__s)); \
+  } while (false)
+
+#define VELOX_RETURN_UNEXPECTED(expected)                    \
+  do {                                                       \
+    auto res = (expected);                                   \
+    VELOX_RETURN_UNEXPECTED_IF(res.hasError(), res.error()); \
+  } while (false)
+
 namespace internal {
 
 /// Common API for extracting Status from either Status or Result<T> (the latter

diff --git a/velox/common/base/tests/StatusTest.cpp b/velox/common/base/tests/StatusTest.cpp
@@ -118,7 +118,7 @@ Status returnNotOk(Status s) {
   return Status::Invalid("invalid");
 }
 
-TEST(StatusTest, macros) {
+TEST(StatusTest, stausMacros) {
   ASSERT_EQ(returnIf(true), Status::Invalid("error"));
   ASSERT_EQ(returnIf(false), Status::OK());
 
@@ -154,5 +154,46 @@ TEST(StatusTest, expected) {
   EXPECT_EQ(result.error(), Status::UserError("division by zero"));
 }
 
+Expected<bool> returnUnexpectedIf(bool cond) {
+  VELOX_RETURN_UNEXPECTED_IF(cond, Status::Invalid("error"));
+  return true;
+}
+
+Expected<bool> returnUnexpected(const Expected<int>& expected) {
+  VELOX_RETURN_UNEXPECTED(expected);
+  return expected.value() == 0;
+}
+
+Expected<bool> returnUnexpectedNotOk(Status s) {
+  VELOX_RETURN_UNEXPECTED_NOT_OK(s);
+  return true;
+}
+
+TEST(StatusTest, expectedMacros) {
+  auto result = returnUnexpectedIf(true);
+  EXPECT_TRUE(result.hasError());
+  EXPECT_EQ(result.error(), Status::Invalid("error"));
+
+  result = returnUnexpectedIf(false);
+  EXPECT_TRUE(result.hasValue());
+  EXPECT_TRUE(result.value());
+
+  result = returnUnexpected(modulo(10, 0));
+  EXPECT_TRUE(result.hasError());
+  EXPECT_EQ(result.error(), Status::UserError("division by zero"));
+
+  result = returnUnexpected(modulo(10, 3));
+  EXPECT_TRUE(result.hasValue());
+  EXPECT_FALSE(result.value());
+
+  result = returnUnexpectedNotOk(Status::UserError("user"));
+  EXPECT_TRUE(result.hasError());
+  EXPECT_EQ(result.error(), Status::UserError("user"));
+
+  result = returnUnexpectedNotOk(Status::OK());
+  EXPECT_TRUE(result.hasValue());
+  EXPECT_TRUE(result.value());
+}
+
 } // namespace
 } // namespace facebook::velox::test