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