diff --git a/cluster/distributed.go b/cluster/distributed.go index 0f0d416..4d779ab 100644 --- a/cluster/distributed.go +++ b/cluster/distributed.go @@ -150,6 +150,15 @@ func (dvm *DistributedVM) SendMessage(message RuntimeMessage) error { // routeMessageToNode sends a message to another node for delivery to the target actor func (dvm *DistributedVM) routeMessageToNode(actorID string, message RuntimeMessage) error { + hops := 0 + if mp, ok := message.(*MessagePayload); ok { + hops = mp.Hops + } + 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, @@ -157,6 +166,7 @@ func (dvm *DistributedVM) routeMessageToNode(actorID string, message RuntimeMess Payload: MessagePayload{ TargetActorID: actorID, Type: message.GetType(), + Hops: hops + 1, }, Timestamp: time.Now(), } @@ -220,6 +230,11 @@ func (dvm *DistributedVM) handleClusterMessage(msg *nats.Msg) { return } + if message.Hops >= MaxRouteHops { + dvm.cluster.logger.Printf("Dropping message for actor %s: exceeded max hops (%d)", message.TargetActorID, MaxRouteHops) + return + } + targetActor := message.TargetActorID if dvm.IsLocalActor(targetActor) { dvm.localRuntime.SendMessage(&message) @@ -236,6 +251,15 @@ func (dvm *DistributedVM) handleClusterMessage(msg *nats.Msg) { // handleRebalanceRequest processes shard rebalancing requests func (dvm *DistributedVM) handleRebalanceRequest(msg ClusterMessage) { + if msg.From == dvm.nodeID { + return + } + + if !dvm.cluster.IsLeader() { + dvm.cluster.logger.Printf("Ignoring rebalance request: not the leader") + return + } + payloadBytes, err := json.Marshal(msg.Payload) if err != nil { dvm.cluster.logger.Printf("Failed to marshal rebalance payload: %v", err) @@ -251,6 +275,10 @@ func (dvm *DistributedVM) handleRebalanceRequest(msg ClusterMessage) { 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)", diff --git a/cluster/manager.go b/cluster/manager.go index 70890c6..ed087e6 100644 --- a/cluster/manager.go +++ b/cluster/manager.go @@ -156,6 +156,8 @@ func (cm *ClusterManager) handleClusterMessage(msg *nats.Msg) { } case "shard_map": cm.handleShardMapUpdate(clusterMsg) + case "migration_update": + cm.handleMigrationUpdate(clusterMsg) default: cm.logger.Printf("Unknown cluster message type: %s", clusterMsg.Type) } @@ -249,6 +251,11 @@ func (cm *ClusterManager) handleRebalanceRequest(msg ClusterMessage) { cm.shardMap = newShardMap cm.mutex.Unlock() + cm.hashRing = NewConsistentHashRing() + for nodeID := range activeNodes { + cm.hashRing.AddNode(nodeID) + } + cm.broadcastShardMap(newShardMap) } @@ -273,6 +280,31 @@ func (cm *ClusterManager) handleMigrationRequest(msg ClusterMessage) { 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) } } @@ -309,6 +341,11 @@ func (cm *ClusterManager) triggerShardRebalancing(reason string) { cm.shardMap = newShardMap cm.mutex.Unlock() + cm.hashRing = NewConsistentHashRing() + for nodeID := range activeNodes { + cm.hashRing.AddNode(nodeID) + } + cm.broadcastShardMap(newShardMap) } @@ -396,6 +433,10 @@ func (cm *ClusterManager) handleShardMapUpdate(msg ClusterMessage) { 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)", @@ -410,12 +451,22 @@ func (cm *ClusterManager) GetShardMap() *ShardMap { 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 @@ -440,3 +491,20 @@ func (cm *ClusterManager) broadcastShardMap(newShardMap *ShardMap) { 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/types.go b/cluster/types.go index 112e555..ee7630e 100644 --- a/cluster/types.go +++ b/cluster/types.go @@ -191,11 +191,15 @@ 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"` + Hops int `json:"hops,omitempty"` } // GetTargetActorID implements RuntimeMessage