Improve our Citus support for modern Citus compatibility. (#939)

* Improve our Citus support for modern Citus compatibility.

  1. Allow rebalance operations with wal_level = 'logical'
  2. Automatically register the coordinator node to itself

* Skip coordinator registration when --skip-pg-hba is used.

There is a chicken-and-egg problem that prevents connecting to the
coordinator this early in the setup of the database. Also, when manual HBA
editing is wanted, it doesn't seem too far off to ask to also take care of
manually running `pg_autoctl activate` on the primary coordinator.

Author: DimCitus

src/bin/pg_autoctl/fsm.c

@@ -201,6 +201,10 @@ KeeperFSMTransition KeeperFSM[] = {
 	/*
 	 * Started as a single, no nothing
 	 */
+	{ INIT_STATE, SINGLE_STATE, NODE_KIND_CITUS_COORDINATOR,
+	  COMMENT_INIT_TO_SINGLE,
+	  &fsm_citus_coordinator_init_primary },
+
 	{ INIT_STATE, SINGLE_STATE, NODE_KIND_CITUS_WORKER,
 	  COMMENT_INIT_TO_SINGLE,
 	  &fsm_citus_worker_init_primary },

src/bin/pg_autoctl/fsm.h

@@ -77,6 +77,7 @@ bool fsm_init_from_standby(Keeper *keeper);
 bool fsm_drop_node(Keeper *keeper);
 
 /* src/bin/pg_autoctl/fsm_transition_citus.c */
+bool fsm_citus_coordinator_init_primary(Keeper *keeper);
 bool fsm_citus_worker_init_primary(Keeper *keeper);
 bool fsm_citus_worker_resume_as_primary(Keeper *keeper);
 bool fsm_citus_coordinator_promote_standby_to_single(Keeper *keeper);

src/bin/pg_autoctl/fsm_transition_citus.c

@@ -41,6 +41,91 @@
 static bool ensure_hostname_is_current_on_coordinator(Keeper *keeper);
 
 
+/*
+ * fsm_citus_coordinator_init_primary initializes a primary coordinator node in
+ * a Citus formation. After doing the usual initialization steps as per the
+ * non-citus version of the FSM, the coordinator node registers itself to the
+ * Citus nodes metadata.
+ */
+bool
+fsm_citus_coordinator_init_primary(Keeper *keeper)
+{
+	KeeperConfig *config = &(keeper->config);
+
+	CoordinatorNodeAddress coordinatorNodeAddress = { 0 };
+	Coordinator coordinator = { 0 };
+	int nodeid = -1;
+
+	if (!fsm_init_primary(keeper))
+	{
+		/* errors have already been logged */
+		return false;
+	}
+
+	/*
+	 * Only Citus workers have more work to do, coordinator are ok. To add
+	 * coordinator to the metadata, users can call "activate" subcommand
+	 * for the coordinator.
+	 */
+	if (keeper->postgres.pgKind != NODE_KIND_CITUS_COORDINATOR)
+	{
+		log_error("BUG: fsm_citus_coordinator_init_primary called for "
+				  "node kind %s",
+				  nodeKindToString(keeper->postgres.pgKind));
+		return false;
+	}
+
+	PostgresSetup *pgSetup = &(config->pgSetup);
+
+	if (pgSetup->hbaLevel <= HBA_EDIT_SKIP)
+	{
+		log_info(
+			"Skipping coordinator registration to itself when --skip-pg-hba "
+			"is used, because we can't connect at pg_autoctl create node time");
+		return true;
+	}
+
+	/*
+	 * We now have a coordinator to talk to: add ourselves as inactive.
+	 */
+	coordinatorNodeAddress.node.port = keeper->config.pgSetup.pgport;
+
+	strlcpy(coordinatorNodeAddress.node.name,
+			keeper->config.name,
+			sizeof(coordinatorNodeAddress.node.name));
+
+	strlcpy(coordinatorNodeAddress.node.host,
+			keeper->config.hostname,
+			sizeof(coordinatorNodeAddress.node.host));
+
+	if (!coordinator_init(&coordinator, &(coordinatorNodeAddress.node), keeper))
+	{
+		log_fatal("Failed to contact the coordinator because its URL is invalid, "
+				  "see above for details");
+		return false;
+	}
+
+	if (!coordinator_add_node(&coordinator, keeper, &nodeid))
+	{
+		/*
+		 * master_add_inactive_node() is idempotent: if the node already has
+		 * been added, nothing changes, in particular if the node is active
+		 * already then the function happily let the node active.
+		 */
+		log_fatal("Failed to add current node to the Citus coordinator, "
+				  "see above for details");
+		return false;
+	}
+
+	log_info("Added coordinator node %s:%d in formation \"%s\" to itself",
+			 keeper->config.hostname,
+			 keeper->config.pgSetup.pgport,
+			 config->formation);
+
+	return true;
+}
+
+
 /*
  * fsm_citus_init_primary initializes a primary worker node in a Citus
  * formation. After doing the usual initialization steps as per the non-citus

src/bin/pg_autoctl/primary_standby.c

@@ -86,6 +86,7 @@ GUC postgres_default_settings_13[] = {
 GUC citus_default_settings_pre_13[] = {
 	DEFAULT_GUC_SETTINGS_FOR_PG_AUTO_FAILOVER_PRE_13,
 	{ "shared_preload_libraries", "'citus,pg_stat_statements'" },
+	{ "wal_level", "logical" },
 	{ "citus.node_conninfo", "'sslmode=prefer'" },
 	{ "citus.cluster_name", "'default'" },
 	{ "citus.use_secondary_nodes", "'never'" },
@@ -96,6 +97,7 @@ GUC citus_default_settings_pre_13[] = {
 GUC citus_default_settings_13[] = {
 	DEFAULT_GUC_SETTINGS_FOR_PG_AUTO_FAILOVER_13,
 	{ "shared_preload_libraries", "'citus,pg_stat_statements'" },
+	{ "wal_level", "logical" },
 	{ "citus.node_conninfo", "'sslmode=prefer'" },
 	{ "citus.cluster_name", "'default'" },
 	{ "citus.use_secondary_nodes", "'never'" },

tests/test_basic_citus_operation.py

@@ -4,6 +4,7 @@
 import os.path
 import time
 import subprocess
+import pprint
 
 cluster = None
 monitor = None
@@ -65,6 +66,12 @@ def test_001_init_coordinator():
     assert coordinator1a.wait_until_state(target_state="primary")
     assert coordinator1b.wait_until_state(target_state="secondary")
 
+    sql = "select nodename, nodeport from pg_dist_node"
+    nodes = coordinator1a.run_sql_query(sql)
+
+    assert nodes[0][0] == str(coordinator1a.vnode.address)
+    assert nodes[0][1] == coordinator1a.port
+
 
 def test_002_init_workers():
     global worker1a

tests/test_citus_skip_pg_hba.py

@@ -52,6 +52,7 @@ def test_001a_init_coordinator():
         # run_sql_query() will need
         # host "172.27.1.1", user "docker", database "postgres"
         hba.write("host postgres docker %s trust\n" % cluster.networkSubnet)
+        hba.write("host citus docker %s trust\n" % cluster.networkSubnet)
         hba.write("host all all %s trust\n" % cluster.networkSubnet)
         hba.write("host replication all %s trust\n" % cluster.networkSubnet)