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

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)
 }

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
 }

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)
+			}
+		}
 	}
 }
 

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