Handle signals properly.

SIGTERM and SIGINT now break the main loop and make the program quit cleanly.
SIGUSR1 triggers a reload of lists (only a stub at the moment).

1 files changed, 78 insertions, 5 deletions

diff --git a/src/main.c b/src/main.c
index fc377e3..e5d5d6a 100644
--- a/src/main.c
+++ b/src/main.c
@@ -18,6 +18,9 @@ Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301  USA
 
 #include "shared.h"
 #include <time.h>
+#include <signal.h>
+#include <sys/select.h>
+#include <fcntl.h>
 
 int server_sockfd;
 
@@ -122,17 +125,51 @@ out:
 	pthread_mutex_unlock(&entry->mutex);
 }
 
+static bool signals_terminate = false;
+static bool signals_reload = false;
+
+// NOTE: The code assumes that signal_handler will only be called in the main
+//       thread, which (as far as I know) is always the case on Linux.
+void signal_handler(int signum)
+{
+	if (signum == SIGINT || signum == SIGTERM)
+		signals_terminate = true;
+
+	if (signum == SIGUSR1)
+		signals_reload = true;
+}
+
+int handle_signals(void)
+{
+	if (signals_terminate) {
+		DEBUG("received a terminating signal\n");
+		return 1;
+	}
+
+	if (signals_reload) {
+		signals_reload = false;
+		DEBUG("reload the lists (not yet implemented)\n");
+	}
+
+	return 0;
+}
+
 #define NUM_WORKERS 8 // FIXME: shouldn't be hardcoded
 
 int main(void)
 {
-	int error = 0; // return value
+	int error = 0 /* the exit code */, flags;
+	struct sigaction sigact = {.sa_handler = signal_handler};
 	struct sockaddr_in sockaddr;
 	pthread_t workers[NUM_WORKERS];
 	size_t i;
 
 	clock_gettime(CLOCK_MONOTONIC_RAW, &time_ref);
 
+	sigaction(SIGINT, &sigact, NULL);
+	sigaction(SIGTERM, &sigact, NULL);
+	sigaction(SIGUSR1, &sigact, NULL);
+
 	if (lists_load("schachts.list", 0)) {
 		eprintf("fatal error: couldn't load the lists\n");
 		goto error_lists;
@@ -145,6 +182,10 @@ int main(void)
 		goto error_socket;
 	}
 
+	// FIXME: can fcntl actually fail (in practice)?
+	flags = fcntl(server_sockfd, F_GETFL, 0);
+	fcntl(server_sockfd, F_SETFL, flags | O_NONBLOCK);
+
 	sockaddr.sin_family = AF_INET;
 	sockaddr.sin_addr.s_addr = INADDR_ANY;
 	sockaddr.sin_port = htons(1337);
@@ -159,17 +200,52 @@ int main(void)
 		pthread_create(workers + i, NULL, worker_main, NULL);
 
 	while (1) {
+		fd_set readfds;
+		int rv;
 		char buffer[256];
 		struct sockaddr_in addr;
 		socklen_t addrlen = sizeof(addr);
 		ssize_t size;
 		uint32_t query;
 
+		FD_ZERO(&readfds);
+		FD_SET(server_sockfd, &readfds);
+
+		// In case the process catches a signal before the main loop
+		// or outside the select call.
+		if (signals_terminate)
+			break;
+
+		// The socket is non-blocking, so wait for incoming data with
+		// select (not recvfrom). Furthermore, select will return if
+		// gets interrupted by a signal, which is main loop breaking
+		// can be easily implemented.
+		rv = select(1, &readfds, NULL, NULL, NULL);
+		if (rv == -1 && errno == EINTR) {
+		interrupted_select:
+			if (handle_signals())
+				break;
+			continue;
+		} else if (rv == -1 && errno != EINTR) {
+			perror("select");
+			break;
+		}
+
 		size = recvfrom(server_sockfd, buffer, sizeof(buffer), 0,
 		                (void*)&addr, &addrlen);
 		if (size == -1) {
+			// Sometimes select can wake up even if there's no
+			// data to read.
+			if (errno == EWOULDBLOCK ||
+			    errno == EAGAIN)
+				continue;
+
+			// Just in case.
+			if (errno == EINTR)
+				goto interrupted_select;
+
 			perror("recvfrom");
-			goto error_recvfrom;
+			break;
 		}
 
 		if ((ntohl(addr.sin_addr.s_addr) & LOCALHOST_MASK)
@@ -190,14 +266,11 @@ int main(void)
 	}
 
 	job_quit();
-
 	for (i = 0; i < NUM_WORKERS; i++)
 		pthread_join(workers[i], NULL);
 
 	db_destroy();
-	// todo: clear all jobs
 
-error_recvfrom:
 error_bind:
 	shutdown(server_sockfd, SHUT_RDWR);
 	close(server_sockfd);