Shardmaster bgw subscribed to new cmds.

Also,
* process_shared_preload_libraries_in_progress check
* shardmaster started only on master
* checking that extension is installed
* handling sigterm on exit and sigusr1 for later cmd canceling

Author: arssher

pg_shardman--0.0.1.sql

@@ -19,7 +19,7 @@ BEGIN
 	RETURN NULL;
 END
 $$ LANGUAGE plpgsql;
-CREATE TRIGGER cmd_log_update
+CREATE TRIGGER cmd_log_inserts
 	AFTER INSERT ON cmd_log
 	FOR EACH STATEMENT
 	EXECUTE PROCEDURE notify_shardmaster();

pg_shardman.control

@@ -4,4 +4,5 @@ default_version = '0.0.1'
 # * learn the scheme when connecting from another node
 # * somehow avoid @extschema@ at init sql at all
 relocatable = false
-schema = shardman
+schema = shardman
+requires = 'pg_pathman, postgres_fdw'

src/pg_shardman.c

@@ -23,11 +23,10 @@
 #include "utils/snapmgr.h"
 #include "executor/spi.h"
 #include "access/xact.h"
+#include "commands/extension.h"
 #include "libpq-fe.h"
 
 
-
-
 /* ensure that extension won't load against incompatible version of Postgres */
 PG_MODULE_MAGIC;
 
@@ -43,6 +42,15 @@ extern void shardmaster_main(Datum main_arg);
 
 static Cmd *next_cmd(void);
 static void update_cmd_status(int64 id, const char *new_status);
+static PGconn *listen_cmd_log_inserts(void);
+static void wait_notify(PGconn *conn);
+static void shardmaster_sigterm(SIGNAL_ARGS);
+static void shardmaster_sigusr1(SIGNAL_ARGS);
+static void pg_shardman_installed(void);
+
+/* flags set by signal handlers */
+static volatile sig_atomic_t got_sigterm = false;
+static volatile sig_atomic_t got_sigusr1 = false;
 
 /* GUC variables */
 static bool shardman_master = false;
@@ -55,6 +63,13 @@ void
 _PG_init()
 {
 	BackgroundWorker shardmaster_worker;
+
+	if (!process_shared_preload_libraries_in_progress)
+	{
+		ereport(ERROR, (errmsg("pg_shardman can only be loaded via shared_preload_libraries"),
+						errhint("Add pg_shardman to shared_preload_libraries.")));
+	}
+
 	DefineCustomBoolVariable("shardman.master",
 							 "This node is the master?",
 							 NULL,
@@ -76,16 +91,19 @@ _PG_init()
 		NULL, NULL, NULL
 		);
 
-	/* register shardmaster */
-	sprintf(shardmaster_worker.bgw_name, "shardmaster");
-	shardmaster_worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
-		BGWORKER_BACKEND_DATABASE_CONNECTION;
-	shardmaster_worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
-	shardmaster_worker.bgw_restart_time = 1;
-	sprintf(shardmaster_worker.bgw_library_name, "pg_shardman");
-	sprintf(shardmaster_worker.bgw_function_name, "shardmaster_main");
-	shardmaster_worker.bgw_notify_pid = 0;
-	RegisterBackgroundWorker(&shardmaster_worker);
+	if (shardman_master)
+	{
+		/* register shardmaster */
+		sprintf(shardmaster_worker.bgw_name, "shardmaster");
+		shardmaster_worker.bgw_flags = BGWORKER_SHMEM_ACCESS |
+			BGWORKER_BACKEND_DATABASE_CONNECTION;
+		shardmaster_worker.bgw_start_time = BgWorkerStart_RecoveryFinished;
+		shardmaster_worker.bgw_restart_time = 10;
+		sprintf(shardmaster_worker.bgw_library_name, "pg_shardman");
+		sprintf(shardmaster_worker.bgw_function_name, "shardmaster_main");
+		shardmaster_worker.bgw_notify_pid = 0;
+		RegisterBackgroundWorker(&shardmaster_worker);
+	}
 }
 
 /*
@@ -96,34 +114,108 @@ shardmaster_main(Datum main_arg)
 {
 	Cmd *cmd;
 	PGconn     *conn;
-	const char *conninfo;
 	elog(LOG, "Shardmaster started");
 
 	/* Connect to the database to use SPI*/
 	BackgroundWorkerInitializeConnection(shardman_master_dbname, NULL);
+	/* sanity check */
+	pg_shardman_installed();
+
+	/* Establish signal handlers before unblocking signals. */
+	pqsignal(SIGTERM, shardmaster_sigterm);
+	pqsignal(SIGUSR1, shardmaster_sigusr1);
+    /* We're now ready to receive signals */
+    BackgroundWorkerUnblockSignals();
+
+	conn = listen_cmd_log_inserts();
+
+	/* main loop */
+	while (!got_sigterm)
+	{
+		while ((cmd = next_cmd()) != NULL)
+		{
+			update_cmd_status(cmd->id, "in progress");
+			elog(LOG, "Working on command %ld, %s", cmd->id, cmd->cmd_type);
+			update_cmd_status(cmd->id, "success");
+		}
+		wait_notify(conn);
+		if (got_sigusr1)
+		{
+			elog(LOG, "SIGUSR1 arrived, aborting current command");
+			update_cmd_status(cmd->id, "canceled");
+			got_sigusr1 = false;
+		}
+	}
+
+	elog(LOG, "Shardmaster received SIGTERM, exiting");
+	PQfinish(conn);
+	proc_exit(0);
+}
+
+/*
+ * Open libpq connection to our server and start listening to cmd_log inserts
+ * notifications.
+ */
+PGconn *
+listen_cmd_log_inserts(void)
+{
+	PGconn *conn;
+	char *conninfo;
+	PGresult   *res;
 
 	conninfo = psprintf("dbname = %s", shardman_master_dbname);
 	conn = PQconnectdb(conninfo);
+	pfree(conninfo);
 	/* Check to see that the backend connection was successfully made */
 	if (PQstatus(conn) != CONNECTION_OK)
 		elog(FATAL, "Connection to database failed: %s",
 			 PQerrorMessage(conn));
 
-
-	/* pg_usleep(10000000L); */
-	while ((cmd = next_cmd()) != NULL)
+	res = PQexec(conn, "LISTEN shardman_cmd_log_update");
+	if (PQresultStatus(res) != PGRES_COMMAND_OK)
 	{
-		elog(LOG, "Working on command %ld, %s", cmd->id, cmd->cmd_type);
-		update_cmd_status(cmd->id, "success");
+		elog(FATAL, "LISTEN command failed: %s", PQerrorMessage(conn));
 	}
+    PQclear(res);
+
+	return conn;
+}
+
+/*
+ * Wait until NOTIFY or signal arrives. If select is alerted, but there are
+ * no notifcations, we also return.
+ */
+void
+wait_notify(PGconn *conn)
+{
+	int			sock;
+	fd_set		input_mask;
+    PGnotify   *notify;
 
+	sock = PQsocket(conn);
+	if (sock < 0)
+		elog(FATAL, "Couldn't get sock from pgconn");
 
-	/* while (1948) */
-	/* { */
+	FD_ZERO(&input_mask);
+	FD_SET(sock, &input_mask);
 
-	/* } */
-	PQfinish(conn);
-	proc_exit(0);
+	if (select(sock + 1, &input_mask, NULL, NULL, NULL) < 0)
+	{
+		if (errno == EINTR)
+			return; /* signal has arrived */
+		elog(FATAL, "select() failed: %s", strerror(errno));
+	}
+
+	PQconsumeInput(conn);
+	/* eat all notifications at once */
+	while ((notify = PQnotifies(conn)) != NULL)
+	{
+		elog(LOG, "NOTIFY %s received from backend PID %d",
+			 notify->relname, notify->be_pid);
+		PQfreemem(notify);
+	}
+
+	return;
 }
 
 /*
@@ -138,7 +230,6 @@ next_cmd(void)
 	MemoryContext oldcxt = CurrentMemoryContext;
 	int e;
 
-	SetCurrentStatementStartTimestamp();
 	StartTransactionCommand();
 	SPI_connect();
 	PushActiveSnapshot(GetTransactionSnapshot());
@@ -184,21 +275,64 @@ update_cmd_status(int64 id, const char *new_status)
 {
 	char *sql;
 	int e;
-	SetCurrentStatementStartTimestamp();
 	StartTransactionCommand();
 	SPI_connect();
 	PushActiveSnapshot(GetTransactionSnapshot());
 
 	sql = psprintf("update shardman.cmd_log set status = '%s' where id = %ld;",
 				   new_status, id);
 	e = SPI_exec(sql, 0);
+	pfree(sql);
 	if (e < 0)
 	{
 		elog(FATAL, "Stmt failed: %s", sql);
 	}
-	pfree(sql);
 
 	PopActiveSnapshot();
 	SPI_finish();
 	CommitTransactionCommand();
 }
+
+/*
+ * Verify that extension is installed locally. We must be connected to db at
+ * this point
+ */
+static void
+pg_shardman_installed(void)
+{
+	bool installed = true;
+
+	StartTransactionCommand();
+	PushActiveSnapshot(GetTransactionSnapshot());
+	if (get_extension_oid("pg_shardman", true) == InvalidOid)
+	{
+		installed = false;
+		elog(WARNING, "pg_shardman library is preloaded, but extenstion is not created");
+	}
+	PopActiveSnapshot();
+	CommitTransactionCommand();
+
+	/* shardmaster won't run without extension */
+	if (!installed)
+		proc_exit(1);
+}
+
+/*
+ * Signal handler for SIGTERM
+ *		Set a flag to let the main loop to terminate.
+ */
+static void
+shardmaster_sigterm(SIGNAL_ARGS)
+{
+	got_sigterm = true;
+}
+
+/*
+ * Signal handler for SIGUSR1
+ *		Set a flag to let the main loop to terminate.
+ */
+static void
+shardmaster_sigusr1(SIGNAL_ARGS)
+{
+	got_sigusr1 = true;
+}