1 #include <c10/macros/Macros.h> 2 #include <c10/util/Backtrace.h> 3 #include <c10/util/env.h> 4 #include <cstdlib> 5 #include <exception> 6 #include <iostream> 7 #include <mutex> 8 #include <optional> 9 10 namespace c10 { 11 class AbortHandlerHelper { 12 public: getInstance()13 static AbortHandlerHelper& getInstance() { 14 #ifdef _WIN32 15 thread_local 16 #endif // _WIN32 17 static AbortHandlerHelper instance; 18 return instance; 19 } 20 set(std::terminate_handler handler)21 void set(std::terminate_handler handler) { 22 std::lock_guard<std::mutex> lk(mutex); 23 if (!inited) { 24 prev = std::set_terminate(handler); 25 curr = std::get_terminate(); 26 inited = true; 27 } 28 } 29 getPrev()30 std::terminate_handler getPrev() const { 31 return prev; 32 } 33 34 private: 35 std::terminate_handler prev = nullptr; 36 std::terminate_handler curr = nullptr; 37 bool inited = false; 38 std::mutex mutex; 39 AbortHandlerHelper() = default; ~AbortHandlerHelper()40 ~AbortHandlerHelper() { 41 // Only restore the handler if we are the current one 42 if (inited && curr == std::get_terminate()) { 43 std::set_terminate(prev); 44 } 45 } 46 47 public: 48 AbortHandlerHelper(AbortHandlerHelper const&) = delete; 49 void operator=(AbortHandlerHelper const&) = delete; 50 }; 51 52 namespace detail { terminate_handler()53C10_ALWAYS_INLINE void terminate_handler() { 54 std::cout << "Unhandled exception caught in c10/util/AbortHandler.h" << '\n'; 55 auto backtrace = get_backtrace(); 56 std::cout << backtrace << '\n' << std::flush; 57 auto prev_handler = AbortHandlerHelper::getInstance().getPrev(); 58 if (prev_handler) { 59 prev_handler(); 60 } else { 61 std::abort(); 62 } 63 } 64 } // namespace detail 65 set_terminate_handler()66C10_ALWAYS_INLINE void set_terminate_handler() { 67 bool use_custom_terminate = false; 68 // On Windows it is enabled by default based on 69 // https://github.com/pytorch/pytorch/pull/50320#issuecomment-763147062 70 #ifdef _WIN32 71 use_custom_terminate = true; 72 #endif // _WIN32 73 auto result = c10::utils::check_env("TORCH_CUSTOM_TERMINATE"); 74 if (result != std::nullopt) { 75 use_custom_terminate = result.value(); 76 } 77 if (use_custom_terminate) { 78 AbortHandlerHelper::getInstance().set(detail::terminate_handler); 79 } 80 } 81 } // namespace c10 82