set_replevel with random node choosing

Author: arssher

init.sql

@@ -24,7 +24,8 @@ $$;
 
 -- available commands
 CREATE TYPE cmd AS ENUM ('add_node', 'rm_node', 'create_hash_partitions',
-						 'move_part', 'create_replica', 'rebalance');
+						 'move_part', 'create_replica', 'rebalance',
+						 'set_replevel');
 -- command status
 CREATE TYPE cmd_status AS ENUM ('waiting', 'canceled', 'failed', 'in progress',
 								'success', 'done');
@@ -143,6 +144,20 @@ BEGIN
 	RETURN c_id;
 END $$ LANGUAGE plpgsql;
 
+-- Add replicas to partitions of table 'relation' until we reach replevel
+-- replicas for each one. Note that it is pointless to set replevel to more than
+-- number of active workers - 1. Replica deletions is not implemented yet.
+CREATE FUNCTION set_replevel(relation text, replevel int) RETURNS int AS $$
+DECLARE
+	c_id int;
+BEGIN
+	INSERT INTO @extschema@.cmd_log VALUES (DEFAULT, 'set_replevel')
+										   RETURNING id INTO c_id;
+	INSERT INTO @extschema@.cmd_opts VALUES (DEFAULT, c_id, relation);
+	INSERT INTO @extschema@.cmd_opts VALUES (DEFAULT, c_id, replevel);
+	RETURN c_id;
+END $$ LANGUAGE plpgsql STRICT;
+
 -- Internal functions
 
 -- Called on shardmaster bgw start. Add itself to nodes table, set id, create

shard.sql

@@ -631,6 +631,18 @@ CREATE FUNCTION reconstruct_table_attrs(relation regclass)
 -- Other funcs
 ------------------------------------------------------------
 
+-- Drop (locally) all partitions of given table, if they exist
+CREATE OR REPLACE FUNCTION drop_parts(relation text, partitions_count int)
+	RETURNS void as $$
+DECLARE
+	r record;
+BEGIN
+	FOR r IN SELECT part_name
+		FROM shardman.gen_part_names(relation, partitions_count) LOOP
+		EXECUTE format('DROP TABLE IF EXISTS %I;', r.part_name);
+	END LOOP;
+END $$ LANGUAGE plpgsql STRICT;
+
 -- generate one-column table with partition names as 'tablename'_'partnum''suffix'
 CREATE FUNCTION gen_part_names(relation text, partitions_count int,
 							   suffix text DEFAULT '')

src/include/pg_shardman.h

@@ -45,6 +45,12 @@ typedef struct Partition
 	int32 owner;
 } Partition;
 
+typedef struct RepCount
+{
+	char *part_name;
+	int64 count;
+} RepCount;
+
 extern void _PG_init(void);
 extern void shardmaster_main(Datum main_arg);
 extern void check_for_sigterm(void);
@@ -59,5 +65,6 @@ extern int32 get_next_node(const char *part_name, int32 node_id);
 extern int32 get_prev_node(const char *part_name, int32 node_id, bool *part_exists);
 extern char *get_partition_relation(const char *part_name);
 extern Partition *get_parts(const char *relation, uint64 *num_parts);
+extern RepCount *get_repcount(const char *relation, uint64 *num_parts);
 
 #endif							/* PG_SHARDMAN_H */

src/include/shard.h

@@ -7,5 +7,6 @@ extern void create_hash_partitions(Cmd *cmd);
 extern void move_part(Cmd *cmd);
 extern void create_replica(Cmd *cmd);
 extern void rebalance(Cmd *cmd);
+extern void set_replevel(Cmd *cmd);
 
 #endif							/* SHARD_H */

src/pg_shardman.c

@@ -208,6 +208,8 @@ shardmaster_main(Datum main_arg)
 				create_replica(cmd);
 			else if (strcmp(cmd->cmd_type, "rebalance") == 0)
 				rebalance(cmd);
+			else if (strcmp(cmd->cmd_type, "set_replevel") == 0)
+				set_replevel(cmd);
 			else
 				shmn_elog(FATAL, "Unknown cmd type %s", cmd->cmd_type);
 		}
@@ -997,3 +999,44 @@ get_parts(const char *relation, uint64 *num_parts)
 	SPI_EPILOG;
 	return parts;
 }
+
+/*
+ * Calculate how many replicas has each partitions of given relation
+ */
+RepCount *
+get_repcount(const char *relation, uint64 *num_parts)
+{
+	char *sql;
+	bool isnull;
+	RepCount *repcounts;
+	TupleDesc rowdesc;
+	MemoryContext spicxt;
+	MemoryContext oldcxt = CurrentMemoryContext;
+	uint64 i;
+
+	SPI_PROLOG;
+	sql = psprintf( /* allocated in SPI ctxt, freed with ctxt release */
+		"select part_name, count(case when prv is not null then 1 end) from"
+		" shardman.partitions where relation = '%s' group by part_name;",
+		relation);
+
+	if (SPI_execute(sql, true, 0) < 0)
+		shmn_elog(FATAL, "Stmt failed : %s", sql);
+	rowdesc = SPI_tuptable->tupdesc;
+
+	*num_parts = SPI_processed;
+	/* We need to allocate in our ctxt, not spi's */
+	spicxt = MemoryContextSwitchTo(oldcxt);
+	repcounts = palloc(sizeof(RepCount) * (*num_parts));
+	for (i = 0; i < *num_parts; i++)
+	{
+		HeapTuple tuple = SPI_tuptable->vals[i];
+		repcounts[i].part_name = SPI_getvalue(tuple, rowdesc, 1);
+		repcounts[i].count = DatumGetInt64(SPI_getbinval(tuple, rowdesc, 2,
+														 &isnull));
+	}
+	MemoryContextSwitchTo(spicxt);
+
+	SPI_EPILOG;
+	return repcounts;
+}

src/shard.c

@@ -7,6 +7,8 @@
  */
 #include "postgres.h"
 
+#include <time.h>
+
 #include "copypart.h"
 #include "pg_shardman.h"
 #include "shard.h"
@@ -61,8 +63,10 @@ create_hash_partitions(Cmd *cmd)
 	 * we have a deadlock between pathman and pg_dump. pfree'd with ctxt
 	 */
 	sql = psprintf(
-		"begin; select create_hash_partitions('%s', '%s', %d); end;"
+		"begin; select shardman.drop_parts('%s', '%d');"
+		" select create_hash_partitions('%s', '%s', %d); end;"
 		"select shardman.gen_create_table_sql('%s', '%s');",
+		relation, partitions_count,
 		relation, expr, partitions_count,
 		relation, connstr);
 
@@ -197,10 +201,9 @@ rebalance(Cmd *cmd)
 	char *relation = cmd->opts[0];
 	uint64 num_workers;
 	int32 *workers = get_workers(&num_workers);
+	uint64 worker_idx;
 	uint64 num_parts;
 	Partition *parts = get_parts(relation, &num_parts);
-	uint64 i;
-	uint64 worker_idx;
 	uint64 part_idx;
 	CopyPartState **tasks = palloc(sizeof(CopyPartState*) * num_parts);
 
@@ -220,10 +223,6 @@ rebalance(Cmd *cmd)
 		return;
 	}
 
-	shmn_elog(DEBUG1, "parts of %s are:", relation);
-	for (i = 0; i < num_parts; i++)
-		shmn_elog(DEBUG1, "%s on %d", parts[i].part_name, parts[i].owner);
-
 	for (part_idx = 0, worker_idx = 0; part_idx < num_parts; part_idx++)
 	{
 		Partition part = parts[part_idx];
@@ -243,5 +242,85 @@ rebalance(Cmd *cmd)
 		return;
 	}
 
+	shmn_elog(INFO, "Relation %s rebalanced:", relation);
 	update_cmd_status(cmd->id, "done");
 }
+
+/*
+ * Add replicas to parts of given relation until we reach replevel replicas
+ * for each one. Worker nodes are choosen in random manner.
+ */
+void
+set_replevel(Cmd *cmd)
+{
+	char *relation = cmd->opts[0];
+	uint32 replevel = atoi(cmd->opts[1]);
+	uint64 num_workers;
+	int32 *workers = get_workers(&num_workers);
+	uint64 num_parts;
+	RepCount *repcounts = NULL;
+	uint64 part_idx;
+	CopyPartState **tasks = NULL;
+	int ntasks;
+	int i;
+
+	if (num_workers == 0)
+	{
+		elog(WARNING, "Set replevel on table %s failed: no active workers",
+			 relation);
+		update_cmd_status(cmd->id, "failed");
+		return;
+	}
+
+	if (replevel > num_workers - 1)
+	{
+		elog(WARNING, "Set replevel on table %s: using replevel %ld instead of"
+			 "%d as we have only %ld active workers",
+			 relation, num_workers - 1, replevel, num_workers);
+		replevel = num_workers - 1;
+	}
+
+	srand(time(NULL));
+	/*
+	 * We can add only one replica per part in one exec_tasks; loop until
+	 * required number of replicas are reached.
+	 */
+	while (1948)
+	{
+		repcounts = get_repcount(relation, &num_parts);
+		tasks = palloc(sizeof(CopyPartState*) * num_parts);
+		ntasks = 0;
+
+		for (part_idx = 0; part_idx < num_parts; part_idx++)
+		{
+			RepCount rc = repcounts[part_idx];
+			if (rc.count < replevel)
+			{
+				CreateReplicaState *crs = palloc0(sizeof(CreateReplicaState));
+				init_cr_state(crs, rc.part_name, workers[rand() % num_workers]);
+				tasks[ntasks] = (CopyPartState *) crs;
+				ntasks++;
+			}
+		}
+
+		if (ntasks == 0)
+			break;
+
+		exec_tasks(tasks, ntasks);
+		check_for_sigterm();
+		if (got_sigusr1)
+		{
+			cmd_canceled(cmd);
+			return;
+		}
+
+		pfree(repcounts);
+		for (i = 0; i < ntasks; i++)
+			pfree(tasks[i]);
+		pfree(tasks);
+	}
+
+	shmn_elog(INFO, "Relation %s now has at least %d replicas", relation,
+			  replevel);
+	update_cmd_status(cmd->id, "success");
+}