f

apache · May 21, 2024 · 7362126 · 7362126
1 parent fcf28bc
commit 7362126
Show file tree

Hide file tree

Showing 2 changed files with 60 additions and 19 deletions.
diff --git a/be/src/util/once.h b/be/src/util/once.h
@@ -22,6 +22,7 @@
 
 #include <atomic>
 #include <mutex>
+#include <stdexcept>
 
 #include "common/exception.h"
 #include "olap/olap_common.h"
@@ -65,22 +66,36 @@ class DorisCallOnce {
     //     return _status;
     // }
 
-    // If the underlying `once_flag` has yet to be invoked, invokes the provided
-    // lambda and stores its return value. Otherwise, returns the stored Status.
-    // If exception occurs in the function, the call flag is not set, user could call
-    // it again. like std::call_once.
+    // If exception occurs in the function, the call flag is set, if user call
+    // it again, the same exception will be thrown.
+    // It is different from std::call_once. This is because if a method is called once
+    // some internal state is changed, it maybe not called again although exception
+    // occurred.
     template <typename Fn>
     ReturnType call(Fn fn) {
+        // Avoid lock to improve performance
         if (has_called()) {
+            if (_eptr) {
+                std::rethrow_exception(_eptr);
+            }
             return _status;
         }
         std::lock_guard l(_flag_lock);
-        // After lock, should check again. Another thread may call
-        // at the same time during first time check has called.
+        // should check again because maybe another thread call successfully.
         if (has_called()) {
+            if (_eptr) {
+                std::rethrow_exception(_eptr);
+            }
             return _status;
         }
-        _status = fn();
+        try {
+            _status = fn();
+        } catch (...) {
+            // Save the exception for next call.
+            _eptr = std::current_exception();
+            _has_called.store(true, std::memory_order_release);
+            std::rethrow_exception(_eptr);
+        }
         _has_called.store(true, std::memory_order_release);
         return _status;
     }
@@ -94,12 +109,19 @@ class DorisCallOnce {
     }
 
     // Return the stored result. The result is only meaningful when `has_called() == true`.
-    ReturnType stored_result() const { return _status; }
+    ReturnType stored_result() const {
+        std::lock_guard l(_flag_lock);
+        if (_eptr) {
+            std::rethrow_exception(_eptr);
+        }
+        return _status;
+    }
 
 private:
     std::atomic<bool> _has_called;
     // std::once_flag _once_flag;
     std::mutex _flag_lock;
+    std::exception_ptr _eptr;
     ReturnType _status;
 };
 

diff --git a/be/test/util/doris_callonce_test.cpp b/be/test/util/doris_callonce_test.cpp
@@ -67,6 +67,7 @@ TEST_F(DorisCallOnceTest, TestExceptionHappens) {
     DorisCallOnce<Status> call1;
     EXPECT_EQ(call1.has_called(), false);
     bool exception_occured = false;
+    bool runtime_occured = false;
     try {
         Status st = call1.call([&]() -> Status {
             throw std::exception();
@@ -77,13 +78,38 @@ TEST_F(DorisCallOnceTest, TestExceptionHappens) {
         exception_occured = true;
     }
     EXPECT_EQ(exception_occured, true);
+    EXPECT_EQ(call1.has_called(), true);
+
+    // call again, should throw the same exception.
+    exception_occured = false;
+    runtime_occured = false;
+    try {
+        Status st = call1.call([&]() -> Status { return Status::InternalError(""); });
+    } catch (...) {
+        // Exception has to throw to the call method
+        exception_occured = true;
+    }
+    EXPECT_EQ(exception_occured, true);
+    EXPECT_EQ(call1.has_called(), true);
+
+    // If call get result, should catch exception.
+    exception_occured = false;
+    runtime_occured = false;
+    try {
+        Status st = call1.stored_result();
+    } catch (...) {
+        // Exception has to throw to the call method
+        exception_occured = true;
+    }
+    EXPECT_EQ(exception_occured, true);
 
     // Test the exception should actually the same one throwed by the callback method.
+    DorisCallOnce<Status> call2;
     exception_occured = false;
-    bool runtime_occured = false;
+    runtime_occured = false;
     try {
         try {
-            Status st = call1.call([&]() -> Status {
+            Status st = call2.call([&]() -> Status {
                 throw std::runtime_error("runtime error happens");
                 return Status::InternalError("");
             });
@@ -100,11 +126,12 @@ TEST_F(DorisCallOnceTest, TestExceptionHappens) {
     EXPECT_EQ(exception_occured, false);
 
     // Test the exception should actually the same one throwed by the callback method.
+    DorisCallOnce<Status> call3;
     exception_occured = false;
     runtime_occured = false;
     try {
         try {
-            Status st = call1.call([&]() -> Status {
+            Status st = call3.call([&]() -> Status {
                 throw std::exception();
                 return Status::InternalError("");
             });
@@ -120,14 +147,6 @@ TEST_F(DorisCallOnceTest, TestExceptionHappens) {
     }
     EXPECT_EQ(runtime_occured, false);
     EXPECT_EQ(exception_occured, true);
-
-    EXPECT_EQ(call1.has_called(), false);
-    EXPECT_EQ(call1.stored_result().code(), ErrorCode::OK);
-
-    Status st = call1.call([&]() -> Status { return Status::InternalError(""); });
-    EXPECT_EQ(call1.has_called(), true);
-    // The error code should not changed
-    EXPECT_EQ(call1.stored_result().code(), ErrorCode::INTERNAL_ERROR);
 }
 
 } // namespace doris