diff --git a/cluster/distributed.go b/cluster/distributed.go index 6e9eb4f..68e7d14 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,43 @@ 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 { + hops := 0 + var body map[string]interface{} + if mp, ok := message.(*MessagePayload); ok { + hops = mp.Hops + body = mp.Body + } + if hops >= MaxRouteHops { + dvm.cluster.logger.Printf("Dropping message for actor %s: exceeded max hops (%d)", actorID, MaxRouteHops) + return fmt.Errorf("message exceeded max hops") + } + + msg := ClusterMessage{ + Type: "route_message", + From: dvm.nodeID, + To: actorID, + Payload: MessagePayload{ + TargetActorID: actorID, + Type: message.GetType(), + Hops: hops + 1, + Body: body, + }, + Timestamp: time.Now(), + } + + return dvm.publishClusterMessage(msg) } // GetActorNode determines which node should handle a specific actor @@ -189,8 +222,10 @@ func (dvm *DistributedVM) handleClusterMessage(msg *nats.Msg) { dvm.localRuntime.LoadModel(&model) case "route_message": - // Handle message routing from other nodes - // Re-marshal and unmarshal to convert map[string]interface{} to concrete type + if clusterMsg.From == dvm.nodeID { + return + } + payloadBytes, err := json.Marshal(clusterMsg.Payload) if err != nil { return @@ -199,7 +234,24 @@ func (dvm *DistributedVM) handleClusterMessage(msg *nats.Msg) { if err := json.Unmarshal(payloadBytes, &message); err != nil { return } - dvm.localRuntime.SendMessage(&message) + + if message.Hops >= MaxRouteHops { + dvm.cluster.logger.Printf("Dropping message for actor %s: exceeded max hops (%d)", message.TargetActorID, MaxRouteHops) + return + } + + targetActor := message.TargetActorID + msg := &MessagePayload{ + TargetActorID: targetActor, + Type: message.Type, + Hops: message.Hops, + Body: message.Body, + } + if dvm.IsLocalActor(targetActor) { + dvm.localRuntime.SendMessage(msg) + } else { + dvm.routeMessageToNode(targetActor, msg) + } case "rebalance": // Handle shard rebalancing requests @@ -209,8 +261,45 @@ 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 + if msg.From == dvm.nodeID { + return + } + + if !dvm.cluster.IsLeader() { + dvm.cluster.logger.Printf("Ignoring rebalance request: not the leader") + return + } + + if dvm.cluster.shardMap == nil { + dvm.cluster.logger.Printf("Shard map is nil, skipping rebalance") + return + } + + 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.hashRing = NewConsistentHashRing() + for nodeID := range newShardMap.Nodes { + dvm.cluster.hashRing.AddNode(nodeID) + } + 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..ed087e6 100644 --- a/cluster/manager.go +++ b/cluster/manager.go @@ -154,6 +154,10 @@ func (cm *ClusterManager) handleClusterMessage(msg *nats.Msg) { if update, ok := clusterMsg.Payload.(NodeUpdate); ok { cm.handleNodeUpdate(update) } + case "shard_map": + cm.handleShardMapUpdate(clusterMsg) + case "migration_update": + cm.handleMigrationUpdate(clusterMsg) default: cm.logger.Printf("Unknown cluster message type: %s", clusterMsg.Type) } @@ -217,16 +221,91 @@ 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.hashRing = NewConsistentHashRing() + for nodeID := range activeNodes { + cm.hashRing.AddNode(nodeID) + } + + 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) + cm.broadcastMigrationUpdate(migration) + } + + if migration.ToNode == cm.nodeID { + cm.logger.Printf("Actor %s assigned to this node, waiting for state import", migration.ActorID) + } +} + +// broadcastMigrationUpdate propagates migration status updates to the cluster +func (cm *ClusterManager) broadcastMigrationUpdate(migration ActorMigration) { + msg := ClusterMessage{ + Type: "migration_update", + From: cm.nodeID, + To: "broadcast", + Payload: migration, + } + + data, err := json.Marshal(msg) + if err != nil { + cm.logger.Printf("Failed to marshal migration update: %v", err) + return + } + + if err := cm.natsConn.Publish("aether.cluster.migration_update", data); err != nil { + cm.logger.Printf("Failed to publish migration update: %v", err) + } } // triggerShardRebalancing initiates shard rebalancing across the cluster @@ -237,12 +316,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 +330,23 @@ 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.hashRing = NewConsistentHashRing() + for nodeID := range activeNodes { + cm.hashRing.AddNode(nodeID) + } + + cm.broadcastShardMap(newShardMap) } // monitorNodes periodically checks node health and updates @@ -319,16 +412,99 @@ 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.hashRing = NewConsistentHashRing() + for nodeID := range newShardMap.Nodes { + cm.hashRing.AddNode(nodeID) + } + 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() defer cm.mutex.RUnlock() // Return a copy to prevent external mutation - return &ShardMap{ + copy := &ShardMap{ Version: cm.shardMap.Version, Shards: make(map[int][]string), Nodes: make(map[string]NodeInfo), UpdateTime: cm.shardMap.UpdateTime, } + + for shardID, nodes := range cm.shardMap.Shards { + copy.Shards[shardID] = append([]string(nil), nodes...) + } + + for nodeID, nodeInfo := range cm.shardMap.Nodes { + copy.Nodes[nodeID] = nodeInfo + } + + return copy +} + +// 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) +} + +// handleMigrationUpdate processes migration status update messages from other nodes +func (cm *ClusterManager) handleMigrationUpdate(msg ClusterMessage) { + var migration ActorMigration + payloadBytes, err := json.Marshal(msg.Payload) + if err != nil { + cm.logger.Printf("Failed to marshal migration update payload: %v", err) + return + } + if err := json.Unmarshal(payloadBytes, &migration); err != nil { + cm.logger.Printf("Failed to unmarshal migration update: %v", err) + return + } + + cm.logger.Printf("Migration update for actor %s: status=%s (from %s)", + migration.ActorID, migration.Status, msg.From) } diff --git a/cluster/shard.go b/cluster/shard.go index 9038815..17a1d29 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,116 @@ 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) { - return nodeID, nil + ring.AddNode(nodeID) + } + + node := ring.GetNode(actorID) + if node == "" { + sortedNodeIDs := make([]string, 0, len(nodes)) + for nodeID := range nodes { + sortedNodeIDs = append(sortedNodeIDs, nodeID) } - i++ + sort.Strings(sortedNodeIDs) + return sortedNodeIDs[0], nil } - // 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 == "" { + sortedNodeIDs := make([]string, 0, len(nodes)) + for nodeID := range nodes { + sortedNodeIDs = append(sortedNodeIDs, nodeID) + } + sort.Strings(sortedNodeIDs) + primaryNode = sortedNodeIDs[0] + } + + 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) + } + } } } diff --git a/cluster/types.go b/cluster/types.go index 112e555..a10af78 100644 --- a/cluster/types.go +++ b/cluster/types.go @@ -191,11 +191,16 @@ func (m *ModelPayload) GetID() string { return m.ID } // GetName implements RuntimeModel func (m *ModelPayload) GetName() string { return m.Name } +// MaxRouteHops is the maximum number of hops a routed message can take before being dropped +const MaxRouteHops = 10 + // MessagePayload is a concrete type for JSON-unmarshaling RuntimeMessage payloads. // Use this when receiving message data over the network. type MessagePayload struct { - TargetActorID string `json:"targetActorId"` - Type string `json:"type"` + TargetActorID string `json:"targetActorId"` + Type string `json:"type"` + Hops int `json:"hops,omitempty"` + Body map[string]interface{} `json:"body,omitempty"` } // GetTargetActorID implements RuntimeMessage