signal_handler.h 3.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106
  1. #pragma once
  2. #include <atomic>
  3. #include <csignal>
  4. #include <mutex>
  5. #include <c10/macros/Export.h>
  6. #include <c10/util/Logging.h>
  7. #if defined(__APPLE__)
  8. #define C10_SUPPORTS_SIGNAL_HANDLER
  9. #elif defined(__linux__) && !defined(C10_DISABLE_SIGNAL_HANDLERS)
  10. #define C10_SUPPORTS_FATAL_SIGNAL_HANDLERS
  11. #define C10_SUPPORTS_SIGNAL_HANDLER
  12. #endif
  13. #if defined(C10_SUPPORTS_FATAL_SIGNAL_HANDLERS)
  14. #include <pthread.h>
  15. #endif
  16. namespace c10 {
  17. class C10_API SignalHandler {
  18. public:
  19. enum class Action { NONE, STOP };
  20. // Constructor. Specify what action to take when a signal is received.
  21. SignalHandler(Action SIGINT_action, Action SIGHUP_action);
  22. ~SignalHandler();
  23. Action CheckForSignals();
  24. bool GotSIGINT();
  25. bool GotSIGHUP();
  26. Action SIGINT_action_;
  27. Action SIGHUP_action_;
  28. unsigned long my_sigint_count_;
  29. unsigned long my_sighup_count_;
  30. };
  31. #if defined(C10_SUPPORTS_FATAL_SIGNAL_HANDLERS)
  32. class C10_API FatalSignalHandler {
  33. // This works by setting up certain fatal signal handlers. Previous fatal
  34. // signal handlers will still be called when the signal is raised. Defaults
  35. // to being off.
  36. public:
  37. C10_API void setPrintStackTracesOnFatalSignal(bool print);
  38. C10_API bool printStackTracesOnFatalSignal();
  39. static FatalSignalHandler& getInstance();
  40. virtual ~FatalSignalHandler();
  41. protected:
  42. explicit FatalSignalHandler();
  43. private:
  44. void installFatalSignalHandlers();
  45. void uninstallFatalSignalHandlers();
  46. static void fatalSignalHandlerStatic(int signum);
  47. void fatalSignalHandler(int signum);
  48. virtual void fatalSignalHandlerPostProcess();
  49. struct sigaction* getPreviousSigaction(int signum);
  50. const char* getSignalName(int signum);
  51. void callPreviousSignalHandler(
  52. struct sigaction* action,
  53. int signum,
  54. siginfo_t* info,
  55. void* ctx);
  56. void stacktraceSignalHandler(bool needsLock);
  57. static void stacktraceSignalHandlerStatic(
  58. int signum,
  59. siginfo_t* info,
  60. void* ctx);
  61. void stacktraceSignalHandler(int signum, siginfo_t* info, void* ctx);
  62. // The mutex protects the bool.
  63. std::mutex fatalSignalHandlersInstallationMutex;
  64. bool fatalSignalHandlersInstalled;
  65. // We need to hold a reference to call the previous SIGUSR2 handler in case
  66. // we didn't signal it
  67. struct sigaction previousSigusr2 {};
  68. // Flag dictating whether the SIGUSR2 handler falls back to previous handlers
  69. // or is intercepted in order to print a stack trace.
  70. std::atomic<bool> fatalSignalReceived;
  71. // Global state set when a fatal signal is received so that backtracing
  72. // threads know why they're printing a stacktrace.
  73. const char* fatalSignalName;
  74. int fatalSignum = -1;
  75. // This wait condition is used to wait for other threads to finish writing
  76. // their stack trace when in fatal sig handler (we can't use pthread_join
  77. // because there's no way to convert from a tid to a pthread_t).
  78. pthread_cond_t writingCond;
  79. pthread_mutex_t writingMutex;
  80. struct signal_handler {
  81. const char* name;
  82. int signum;
  83. struct sigaction previous;
  84. };
  85. static signal_handler kSignalHandlers[];
  86. };
  87. #endif // defined(C10_SUPPORTS_SIGNAL_HANDLER)
  88. } // namespace c10