| 123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081 |
- #include <c10/macros/Macros.h>
- #include <c10/util/Backtrace.h>
- #include <c10/util/env.h>
- #include <cstdlib>
- #include <exception>
- #include <iostream>
- #include <mutex>
- #include <optional>
- namespace c10 {
- class AbortHandlerHelper {
- public:
- static AbortHandlerHelper& getInstance() {
- #ifdef _WIN32
- thread_local
- #endif // _WIN32
- static AbortHandlerHelper instance;
- return instance;
- }
- void set(std::terminate_handler handler) {
- std::lock_guard<std::mutex> lk(mutex);
- if (!inited) {
- prev = std::set_terminate(handler);
- curr = std::get_terminate();
- inited = true;
- }
- }
- std::terminate_handler getPrev() const {
- return prev;
- }
- private:
- std::terminate_handler prev = nullptr;
- std::terminate_handler curr = nullptr;
- bool inited = false;
- std::mutex mutex;
- AbortHandlerHelper() = default;
- ~AbortHandlerHelper() {
- // Only restore the handler if we are the current one
- if (inited && curr == std::get_terminate()) {
- std::set_terminate(prev);
- }
- }
- public:
- AbortHandlerHelper(AbortHandlerHelper const&) = delete;
- void operator=(AbortHandlerHelper const&) = delete;
- };
- namespace detail {
- C10_ALWAYS_INLINE void terminate_handler() {
- std::cout << "Unhandled exception caught in c10/util/AbortHandler.h" << '\n';
- auto backtrace = get_backtrace();
- std::cout << backtrace << '\n' << std::flush;
- auto prev_handler = AbortHandlerHelper::getInstance().getPrev();
- if (prev_handler) {
- prev_handler();
- } else {
- std::abort();
- }
- }
- } // namespace detail
- C10_ALWAYS_INLINE void set_terminate_handler() {
- bool use_custom_terminate = false;
- // On Windows it is enabled by default based on
- // https://github.com/pytorch/pytorch/pull/50320#issuecomment-763147062
- #ifdef _WIN32
- use_custom_terminate = true;
- #endif // _WIN32
- auto result = c10::utils::check_env("TORCH_CUSTOM_TERMINATE");
- if (result != std::nullopt) {
- use_custom_terminate = result.value();
- }
- if (use_custom_terminate) {
- AbortHandlerHelper::getInstance().set(detail::terminate_handler);
- }
- }
- } // namespace c10
|