Signal handling revisited.

Macros for easy cmd canceled & sigterm check.

Author: arssher

shard.sql

@@ -220,7 +220,6 @@ BEGIN
 		END IF;
 		-- Drop old table anyway;
 		EXECUTE format('DROP TABLE IF EXISTS %I', NEW.part_name);
-
 	ELSEIF me = NEW.owner THEN -- dst node
 		RAISE DEBUG '[SHARDMAN %] part_moved trigger on dst node',
 			me;

src/copypart.c

@@ -440,7 +440,7 @@ exec_tasks(CopyPartState **tasks, int ntasks)
 	if ((epfd = epoll_create1(0)) == -1)
 		shmn_elog(FATAL, "epoll_create1 failed");
 
-	while (unfinished_tasks > 0 && !got_sigusr1 && !got_sigterm)
+	while (unfinished_tasks > 0 && !signal_pending())
 	{
 		timeout = calc_timeout(&timeout_states);
 		e = epoll_wait(epfd, evlist, MAX_EVENTS, timeout);

src/include/pg_shardman.h

@@ -3,23 +3,48 @@
 
 #include <signal.h>
 
+#include "miscadmin.h"
+
 #define shmn_elog(level,fmt,...) elog(level, "[SHARDMAN] " fmt, ## __VA_ARGS__)
 
 #define SPI_PROLOG do { \
-		StartTransactionCommand(); \
-		SPI_connect(); \
-		PushActiveSnapshot(GetTransactionSnapshot()); \
-	} while (0);
+	StartTransactionCommand(); \
+	SPI_connect(); \
+	PushActiveSnapshot(GetTransactionSnapshot()); \
+} while (0);
 
 #define SPI_EPILOG do { \
-		PopActiveSnapshot(); \
-		SPI_finish(); \
-		CommitTransactionCommand(); \
-	} while (0);
+	PopActiveSnapshot(); \
+	SPI_finish(); \
+	CommitTransactionCommand(); \
+} while (0);
 
 /* flags set by signal handlers */
 extern volatile sig_atomic_t got_sigterm;
 extern volatile sig_atomic_t got_sigusr1;
+/*
+ * Most probably CHECK_FOR_INTERRTUPS here is useless since we handle
+ * SIGTERM ourselves (to get adequate log message) and don't need anything
+ * else, but just to be sure...
+ */
+#define SHMN_CHECK_FOR_INTERRUPTS() \
+do { \
+	check_for_sigterm(); \
+	CHECK_FOR_INTERRUPTS(); \
+} while (0)
+/*
+ * Additionally check for SIGUSR1; if it has arrived, mark cmd as canceled and
+ * return from current function.
+ */
+#define SHMN_CHECK_FOR_INTERRUPTS_CMD(cmd) \
+do { \
+	SHMN_CHECK_FOR_INTERRUPTS(); \
+	if (got_sigusr1) \
+	{ \
+		cmd_canceled((cmd)); \
+		return; \
+	} \
+} while(0)
 
 /* GUC variables */
 extern bool shardman_shardlord;
@@ -53,10 +78,11 @@ typedef struct RepCount
 
 extern void _PG_init(void);
 extern void shardlord_main(Datum main_arg);
+extern bool signal_pending(void);
 extern void check_for_sigterm(void);
+extern void cmd_canceled(Cmd *cmd);
 extern uint64 void_spi(char *sql);
 extern void update_cmd_status(int64 id, const char *new_status);
-extern void cmd_canceled(Cmd *cmd);
 extern char *get_worker_node_connstr(int32 node_id);
 extern int32 *get_workers(uint64 *num_workers);
 extern int32 get_primary_owner(const char *part_name);

src/pg_shardman.c

@@ -431,14 +431,16 @@ shardlord_sigterm(SIGNAL_ARGS)
 
 /*
  * Signal handler for SIGUSR1
- *		Set a flag to let the main loop to terminate.
+ *		Set a flag to let cancel a command
  */
 void
 shardlord_sigusr1(SIGNAL_ARGS)
 {
 	got_sigusr1 = true;
 }
 
+bool signal_pending(void) { return got_sigterm || got_sigusr1; }
+
 /*
  * Cleanup and exit in case of SIGTERM
  */
@@ -488,7 +490,7 @@ add_node(Cmd *cmd)
 
 	shmn_elog(INFO, "Adding node %s", connstr);
 	/* Try to execute command indefinitely until it succeeded or canceled */
-	while (!got_sigusr1 && !got_sigterm)
+	while (1948)
 	{
 		conn = PQconnectdb(connstr);
 		if (PQstatus(conn) != CONNECTION_OK)
@@ -610,10 +612,8 @@ add_node(Cmd *cmd)
 		shmn_elog(LOG, "Attempt to execute add_node failed, sleeping and retrying");
 		/* TODO: sleep using waitlatch? */
 		pg_usleep(shardman_cmd_retry_naptime * 1000L);
+		SHMN_CHECK_FOR_INTERRUPTS_CMD(cmd);
 	}
-	check_for_sigterm();
-
-	cmd_canceled(cmd);
 }
 
 /* See sql func */

src/shard.c

@@ -71,7 +71,7 @@ create_hash_partitions(Cmd *cmd)
 		relation, connstr);
 
 	/* Try to execute command indefinitely until it succeeded or canceled */
-	while (!got_sigusr1 && !got_sigterm)
+	while (1948)
 	{
 		conn = PQconnectdb(connstr);
 		if (PQstatus(conn) != CONNECTION_OK)
@@ -127,22 +127,15 @@ create_hash_partitions(Cmd *cmd)
 				  " sleeping and retrying");
 		/* TODO: sleep using waitlatch? */
 		pg_usleep(shardman_cmd_retry_naptime * 1000L);
+		SHMN_CHECK_FOR_INTERRUPTS_CMD(cmd);
 	}
-	check_for_sigterm();
-
-	cmd_canceled(cmd);
 }
 
 /* Update status of cmd consisting of single task after exec_tasks finishes */
 void
 cmd_single_task_exec_finished(Cmd *cmd, CopyPartState *cps)
 {
-	check_for_sigterm();
-	if (got_sigusr1)
-	{
-		cmd_canceled(cmd);
-		return;
-	}
+	SHMN_CHECK_FOR_INTERRUPTS_CMD(cmd);
 
 	Assert(cps->res != TASK_IN_PROGRESS);
 	if (cps->res == TASK_FAILED)
@@ -235,12 +228,7 @@ rebalance(Cmd *cmd)
 	}
 
 	exec_tasks(tasks, num_parts);
-	check_for_sigterm();
-	if (got_sigusr1)
-	{
-		cmd_canceled(cmd);
-		return;
-	}
+	SHMN_CHECK_FOR_INTERRUPTS_CMD(cmd);
 
 	shmn_elog(INFO, "Relation %s rebalanced:", relation);
 	update_cmd_status(cmd->id, "done");
@@ -307,12 +295,7 @@ set_replevel(Cmd *cmd)
 			break;
 
 		exec_tasks(tasks, ntasks);
-		check_for_sigterm();
-		if (got_sigusr1)
-		{
-			cmd_canceled(cmd);
-			return;
-		}
+		SHMN_CHECK_FOR_INTERRUPTS_CMD(cmd);
 
 		pfree(repcounts);
 		for (i = 0; i < ntasks; i++)