-
Notifications
You must be signed in to change notification settings - Fork 8
Expand file tree
/
Copy pathsignal_handler.cpp
More file actions
137 lines (123 loc) · 3.51 KB
/
signal_handler.cpp
File metadata and controls
137 lines (123 loc) · 3.51 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
#include <boost/bind.hpp>
#include <glog/logging.h>
#include <signal.h>
#include <csignal>
#include "caffe/util/signal_handler.h"
namespace {
static volatile sig_atomic_t got_sigint = false;
static volatile sig_atomic_t got_sighup = false;
static bool already_hooked_up = false;
void handle_signal(int signal) {
switch (signal) {
#ifdef _MSC_VER
case SIGBREAK: // there is no SIGHUP in windows, take SIGBREAK instead.
got_sighup = true;
break;
#else
case SIGHUP:
got_sighup = true;
break;
#endif
case SIGINT:
got_sigint = true;
break;
}
}
void HookupHandler() {
if (already_hooked_up) {
LOG(FATAL) << "Tried to hookup signal handlers more than once.";
}
already_hooked_up = true;
#ifdef _MSC_VER
if (signal(SIGBREAK, handle_signal) == SIG_ERR) {
LOG(FATAL) << "Cannot install SIGBREAK handler.";
}
if (signal(SIGINT, handle_signal) == SIG_ERR) {
LOG(FATAL) << "Cannot install SIGINT handler.";
}
#else
struct sigaction sa;
// Setup the handler
sa.sa_handler = &handle_signal;
// Restart the system call, if at all possible
sa.sa_flags = SA_RESTART;
// Block every signal during the handler
sigfillset(&sa.sa_mask);
// Intercept SIGHUP and SIGINT
if (sigaction(SIGHUP, &sa, NULL) == -1) {
LOG(FATAL) << "Cannot install SIGHUP handler.";
}
if (sigaction(SIGINT, &sa, NULL) == -1) {
LOG(FATAL) << "Cannot install SIGINT handler.";
}
#endif
}
// Set the signal handlers to the default.
void UnhookHandler() {
if (already_hooked_up) {
#ifdef _MSC_VER
if (signal(SIGBREAK, SIG_DFL) == SIG_ERR) {
LOG(FATAL) << "Cannot uninstall SIGBREAK handler.";
}
if (signal(SIGINT, SIG_DFL) == SIG_ERR) {
LOG(FATAL) << "Cannot uninstall SIGINT handler.";
}
#else
struct sigaction sa;
// Setup the sighub handler
sa.sa_handler = SIG_DFL;
// Restart the system call, if at all possible
sa.sa_flags = SA_RESTART;
// Block every signal during the handler
sigfillset(&sa.sa_mask);
// Intercept SIGHUP and SIGINT
if (sigaction(SIGHUP, &sa, NULL) == -1) {
LOG(FATAL) << "Cannot uninstall SIGHUP handler.";
}
if (sigaction(SIGINT, &sa, NULL) == -1) {
LOG(FATAL) << "Cannot uninstall SIGINT handler.";
}
#endif
already_hooked_up = false;
}
}
// Return true iff a SIGINT has been received since the last time this
// function was called.
bool GotSIGINT() {
bool result = got_sigint;
got_sigint = false;
return result;
}
// Return true iff a SIGHUP has been received since the last time this
// function was called.
bool GotSIGHUP() {
bool result = got_sighup;
got_sighup = false;
return result;
}
} // namespace
namespace caffe {
SignalHandler::SignalHandler(SolverAction::Enum SIGINT_action,
SolverAction::Enum SIGHUP_action):
SIGINT_action_(SIGINT_action),
SIGHUP_action_(SIGHUP_action) {
HookupHandler();
}
SignalHandler::~SignalHandler() {
UnhookHandler();
}
SolverAction::Enum SignalHandler::CheckForSignals() const {
if (GotSIGHUP()) {
return SIGHUP_action_;
}
if (GotSIGINT()) {
return SIGINT_action_;
}
return SolverAction::NONE;
}
// Return the function that the solver can use to find out if a snapshot or
// early exit is being requested.
ActionCallback SignalHandler::GetActionFunction() {
return boost::bind(&SignalHandler::CheckForSignals, this);
}
} // namespace caffe