Fix issue with rebalancing never happens with auto-sharding (#531)

* Fix issue with rebalancing never happens with auto-sharding

* Add test for rebalancing scheme generation to make sure it works fine

* Fix various critical issues in auto sharding rebalancing logic

* Fix sharding rebalancing: recreate distributed table not just drop

* Fix issue with cluster script

* Update schema in proper time to fix issue with 5 nodes sharding rebalancing

* Remove wrong added script

* Remove debugging information

* Fix: activate auto sharding on new node appeared in the cluster

* feat(sharding): add shard creation for new nodes during rebalance

- Log distinct messages for rebalancing due to inactive vs new nodes
- Detect new nodes joining cluster and create missing shard tables
- Set up replication from existing nodes to new nodes for new shards
- Improve shard map initialization to avoid undefined keys

* feat(sharding): skip rebalancing when replication factor is 1

- Prevent data movement and replication setup if original RF=1 to avoid data loss
- Only add new nodes without redistributing shards when RF=1
- Perform full rebalancing and replication setup only if RF > 1
- Add method to calculate original replication factor from schema
- Update rebalance logic and shard creation to respect RF for safe operations

* Fix test

* Fix cognitive complexity

Author: donhardman

bin/cluster

@@ -10,8 +10,8 @@ if ((nodes < 2)); then
 	echo "Number of nodes must be greater than 1"
 	exit 1
 fi
-if ((nodes > 4)); then
-	echo "Number of nodes must be less than or equal to 555"
+if ((nodes > 5)); then
+	echo "Number of nodes must be less than or equal to 5"
 	exit 1
 fi
 

src/Plugin/Sharding/Operator.php

@@ -100,11 +100,9 @@ public function checkMaster(): static {
 	public function checkBalance(): static {
 		$cluster = $this->getCluster();
 		$queue = $this->getQueue();
+		// We get inactive nodes to exclude them from the rebalance in case of outages
+		// It may be an empty list if we're adding a new node to the cluster, which is fine
 		$inactiveNodes = $cluster->getInactiveNodes();
-		// Do rebalance in case we have inactive nodes
-		if ($inactiveNodes->isEmpty()) {
-			return $this;
-		}
 
 		// Do exclude repeated rebalancing we take hash of active nodes
 		// and if the same, we do nothing, otherwise, it shows the change
@@ -115,7 +113,12 @@ public function checkBalance(): static {
 		if ($clusterHash === $currentHash) {
 			return $this;
 		}
-		Buddy::info("Rebalancing due to inactive nodes: {$inactiveNodes->join(', ')}");
+
+		if ($inactiveNodes->count() > 0) {
+			Buddy::info("Rebalancing due to inactive nodes: {$inactiveNodes->join(', ')}");
+		} else {
+			Buddy::info('Rebalancing due to new nodes joined');
+		}
 
 		// Get all tables from the state we have
 		$list = $this->state->listRegex('table:.+');

src/Plugin/Sharding/Table.php

@@ -113,14 +113,19 @@ public function getConnectedNodes(Set $shards): Set {
 	 * @return Vector<array{node:string,shards:Set<int>}>
 	 */
 	public function getExternalNodeShards(Set $shards): Vector {
+		// There is a case when we have ONLY distributed table with shards
+		// but all shards are remote agents and nothing locally
+		$shardsCond = $shards->count() > 0
+			? "AND ANY(shards) not in ({$shards->join(',')})"
+			: '';
+
 		$query = "
 		SELECT node, shards FROM {$this->table}
 		WHERE
 		cluster = '{$this->cluster->name}'
 		AND
 		table = '{$this->name}'
-		AND
-		ANY(shards) not in ({$shards->join(',')})
+		$shardsCond
 		ORDER BY id ASC
 		";
 
@@ -192,27 +197,21 @@ public function shard(
 		$nodes = new Set;
 
 		$schema = $this->configureNodeShards($shardCount, $replicationFactor);
-		$reduceFn = function (Map $clusterMap, array $row) use ($queue, $replicationFactor, &$nodes, &$nodeShardsMap) {
+		$reduceFn = function (Map $clusterMap, array $row) use ($queue, &$nodes, &$nodeShardsMap) {
 			/** @var Map<string,Cluster> $clusterMap */
 			$nodes->add($row['node']);
 			$nodeShardsMap[$row['node']] = $row['shards'];
 
 			foreach ($row['shards'] as $shard) {
 				$connectedNodes = $this->getConnectedNodes(new Set([$shard]));
 
-				if ($replicationFactor > 1) {
-					$clusterMap = $this->handleReplication(
-						$row['node'],
-						$queue,
-						$connectedNodes,
-						$clusterMap,
-						$shard
-					);
-				} else {
-					// If no replication, add create table shard SQL to queue
-					$sql = $this->getCreateTableShardSQL($shard);
-					$queue->add($row['node'], $sql);
-				}
+				$clusterMap = $this->handleReplication(
+					$row['node'],
+					$queue,
+					$connectedNodes,
+					$clusterMap,
+					$shard
+				);
 			}
 
 			return $clusterMap;
@@ -227,10 +226,7 @@ public function shard(
 		/** @var Set<int> */
 		$queueIds = new Set;
 		foreach ($nodeShardsMap as $node => $shards) {
-			// Do nothing when no shards present for this node
-			if (!$shards->count()) {
-				continue;
-			}
+			// Even when no shards, we still create distributed table
 			$sql = $this->getCreateShardedTableSQL($shards);
 			$queueId = $queue->add($node, $sql);
 			$queueIds->add($queueId);
@@ -311,6 +307,13 @@ protected function handleReplication(
 		Map $clusterMap,
 		int $shard
 	): Map {
+		// If no replication, add create table shard SQL to queue
+		if ($connectedNodes->count() === 1) {
+			$sql = $this->getCreateTableShardSQL($shard);
+			$queue->add($node, $sql);
+			return $clusterMap;
+		}
+
 		$clusterName = static::getClusterName($connectedNodes);
 		$hasCluster = isset($clusterMap[$clusterName]);
 		if ($hasCluster) {
@@ -354,9 +357,6 @@ public function rebalance(Queue $queue): void {
 			$schema = $this->getShardSchema();
 			$allNodes = $this->cluster->getNodes();
 			$inactiveNodes = $this->cluster->getInactiveNodes();
-			if (!$inactiveNodes->count()) {
-				return;
-			}
 			$activeNodes = $allNodes->diff($inactiveNodes);
 			$newSchema = Util::rebalanceShardingScheme($schema, $activeNodes);
 
@@ -376,6 +376,7 @@ public function rebalance(Queue $queue): void {
 				$clusterMap[$clusterName] = $cluster;
 			}
 
+			// Get affected schema with nodes that are out
 			$affectedSchema = $schema->filter(
 				fn ($row) => $inactiveNodes->contains($row['node'])
 			);
@@ -387,7 +388,7 @@ public function rebalance(Queue $queue): void {
 				// First thing first, remove from inactive node using the queue
 				$this->cleanUpNode($queue, $row['node'], $row['shards'], $processedTables);
 
-				// Do real rebaliance now
+				// Do real rebalancing now
 				foreach ($row['shards'] as $shard) {
 					/** @var Set<string> */
 					$nodesForShard = new Set;
@@ -438,20 +439,34 @@ public function rebalance(Queue $queue): void {
 
 			/** @var Set<int> */
 			$queueIds = new Set;
+			foreach ($clusterMap as $cluster) {
+				$cluster->processPendingTables($queue);
+			}
+
+			// Handle new nodes that need shard creation
+			$originalNodes = new Set($schema->map(fn($row) => $row['node']));
+			$newNodes = $activeNodes->diff($originalNodes);
+
+
+
+			if ($newNodes->count() > 0) {
+				$this->handleShardCreationForRebalancing($queue, $schema, $newSchema, $clusterMap);
+			}
+
+			// At this case we update schema
+			// before creating distributed table
+			$this->updateScheme($newSchema);
 			foreach ($newSchema as $row) {
-				$sql = "DROP TABLE {$this->name} OPTION force=1";
+				// We should drop distributed table everywhere
+				// even when node has ONLY it but may have no shards on it
+				$sql = "DROP TABLE IF EXISTS {$this->name} OPTION force=1";
 				$queueId = $queue->add($row['node'], $sql);
 				$queueIds->add($queueId);
-				// Do nothing when no shards present for this node
-				if (!$row['shards']->count()) {
-					continue;
-				}
+
 				$sql = $this->getCreateShardedTableSQL($row['shards']);
 				$queueId = $queue->add($row['node'], $sql);
 				$queueIds->add($queueId);
 			}
-
-			$this->updateScheme($newSchema);
 		} catch (\Throwable $t) {
 			var_dump($t->getMessage());
 		}
@@ -659,10 +674,12 @@ protected function getCreateShardedTableSQL(Set $shards): string {
 		$map = new Map;
 		foreach ($nodes as $row) {
 			foreach ($row['shards'] as $shard) {
-				$map[$shard] ??= new Set;
+				if (!$map->hasKey($shard)) {
+					$map[$shard] = new Set();
+				}
 				$shardName = $this->getShardName($shard);
-				// @phpstan-ignore-next-line
-				$map[$shard]->add("{$row['node']}:{$shardName}");
+
+				$map[$shard]?->add("{$row['node']}:{$shardName}");
 			}
 		}
 
@@ -738,6 +755,105 @@ public function setup(): void {
 	 * @return Set<int>
 	 */
 	protected static function parseShards(string $shards): Set {
-		return new Set(array_map('intval', explode(',', $shards)));
+		return trim($shards) !== ''
+			? new Set(array_map('intval', explode(',', $shards)))
+			: new Set
+		;
+	}
+
+	/**
+	 * Handle shard creation for rebalancing (all nodes that need new shards)
+	 *
+	 * SAFETY: Respects original replication factor - with RF=1, only creates shard tables
+	 * but does NOT set up replication to prevent data movement and potential data loss.
+	 *
+	 * @param Queue $queue
+	 * @param Vector<array{node:string,shards:Set<int>,connections:Set<string>}> $oldSchema
+	 * @param Vector<array{node:string,shards:Set<int>,connections:Set<string>}> $newSchema
+	 * @param Map<string,Cluster> $clusterMap
+	 * @return void
+	 */
+	protected function handleShardCreationForRebalancing(
+		Queue $queue,
+		Vector $oldSchema,
+		Vector $newSchema,
+		Map $clusterMap
+	): void {
+		// Calculate original replication factor to ensure safe operations
+		$originalRf = $this->calculateReplicationFactor($oldSchema);
+
+		// Create map of old schema for comparison
+		/** @var Map<string,Set<int>> */
+		$oldShardMap = new Map();
+		foreach ($oldSchema as $row) {
+			$oldShardMap[$row['node']] = $row['shards'];
+		}
+
+		foreach ($newSchema as $row) {
+			$oldShards = $oldShardMap->get($row['node'], new Set());
+			$newShards = $row['shards'];
+			$shardsToCreate = $newShards->diff($oldShards);
+
+			if ($shardsToCreate->isEmpty()) {
+				continue;
+			}
+
+			// Create missing shard tables on this node
+			foreach ($shardsToCreate as $shard) {
+				$sql = $this->getCreateTableShardSQL($shard);
+				$queue->add($row['node'], $sql);
+
+				// Find nodes that already have this shard in old schema for replication
+				$existingNodesWithShard = $oldSchema->filter(
+					fn($existingRow) => $existingRow['shards']->contains($shard)
+				);
+
+				// If there are existing nodes with this shard, set up replication
+				// BUT only if original RF > 1 (safe to replicate)
+				if ($existingNodesWithShard->count() <= 0 || $originalRf === 1) {
+					continue;
+				}
+
+				$sourceNode = $existingNodesWithShard->first()['node'];
+				$connectedNodes = new Set([$row['node'], $sourceNode]);
+
+				// Set up cluster replication for this shard
+				$this->handleReplication(
+					$sourceNode,
+					$queue,
+					$connectedNodes,
+					$clusterMap,
+					$shard
+				);
+			}
+		}
+	}
+
+	/**
+	 * Calculate the original replication factor from the schema
+	 * @param Vector<array{node:string,shards:Set<int>,connections:Set<string>}> $schema
+	 * @return int
+	 */
+	private function calculateReplicationFactor(Vector $schema): int {
+		if ($schema->isEmpty()) {
+			return 1;
+		}
+
+		// Count how many nodes have each shard
+		$shardCounts = new Map();
+
+		foreach ($schema as $row) {
+			foreach ($row['shards'] as $shard) {
+				$currentCount = $shardCounts->get($shard, 0);
+				$shardCounts->put($shard, $currentCount + 1);
+			}
+		}
+
+		if ($shardCounts->isEmpty()) {
+			return 1;
+		}
+
+		// The replication factor is the maximum count of any shard
+		return max($shardCounts->values()->toArray());
 	}
 }