AbortHandler.h 2.0 KB

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