aether/cluster/shard.go

package cluster

import (
	"crypto/sha256"
	"encoding/binary"
	"fmt"
	"hash"
	"hash/fnv"
	"sort"
	"time"
)

// MigrationStatus tracks actor migration progress
type MigrationStatus string

const (
	MigrationPending    MigrationStatus = "pending"
	MigrationInProgress MigrationStatus = "in_progress"
	MigrationCompleted  MigrationStatus = "completed"
	MigrationFailed     MigrationStatus = "failed"
)

// PlacementStrategy determines where to place new actors
type PlacementStrategy interface {
	PlaceActor(actorID string, shardMap *ShardMap, nodes map[string]*NodeInfo) (string, error)
	RebalanceShards(shardMap *ShardMap, nodes map[string]*NodeInfo) (*ShardMap, error)
}

// ShardManager handles actor placement and distribution
type ShardManager struct {
	shardCount  int
	shardMap    *ShardMap
	hasher      hash.Hash
	placement   PlacementStrategy
	replication int
}

// NewShardManager creates a new shard manager with default configuration
func NewShardManager(shardCount, replication int) *ShardManager {
	return NewShardManagerWithConfig(ShardConfig{
		ShardCount:        shardCount,
		ReplicationFactor: replication,
	})
}

// NewShardManagerWithConfig creates a new shard manager with custom configuration
func NewShardManagerWithConfig(config ShardConfig) *ShardManager {
	// Apply defaults for zero values
	shardCount := config.ShardCount
	if shardCount == 0 {
		shardCount = DefaultNumShards
	}
	replication := config.ReplicationFactor
	if replication == 0 {
		replication = 1
	}

	return &ShardManager{
		shardCount:  shardCount,
		shardMap:    &ShardMap{Shards: make(map[int][]string), Nodes: make(map[string]NodeInfo)},
		hasher:      fnv.New64a(),
		placement:   &ConsistentHashPlacement{},
		replication: replication,
	}
}

// GetShard returns the shard number for a given actor ID
func (sm *ShardManager) GetShard(actorID string) int {
	h := sha256.Sum256([]byte(actorID))
	shardID := binary.BigEndian.Uint32(h[:4]) % uint32(sm.shardCount)
	return int(shardID)
}

// GetShardNodes returns the nodes responsible for a shard
func (sm *ShardManager) GetShardNodes(shardID int) []string {
	if nodes, exists := sm.shardMap.Shards[shardID]; exists {
		return nodes
	}
	return []string{}
}

// AssignShard assigns a shard to specific nodes
func (sm *ShardManager) AssignShard(shardID int, nodes []string) {
	if sm.shardMap.Shards == nil {
		sm.shardMap.Shards = make(map[int][]string)
	}
	sm.shardMap.Shards[shardID] = nodes
}

// GetPrimaryNode returns the primary node for a shard
func (sm *ShardManager) GetPrimaryNode(shardID int) string {
	nodes := sm.GetShardNodes(shardID)
	if len(nodes) > 0 {
		return nodes[0] // First node is primary
	}
	return ""
}

// GetReplicaNodes returns the replica nodes for a shard
func (sm *ShardManager) GetReplicaNodes(shardID int) []string {
	nodes := sm.GetShardNodes(shardID)
	if len(nodes) > 1 {
		return nodes[1:] // All nodes except first are replicas
	}
	return []string{}
}

// UpdateShardMap updates the entire shard map
func (sm *ShardManager) UpdateShardMap(newShardMap *ShardMap) {
	sm.shardMap = newShardMap
}

// GetShardMap returns a copy of the current shard map
func (sm *ShardManager) GetShardMap() *ShardMap {
	// Return a deep copy to prevent external mutation
	copy := &ShardMap{
		Version:    sm.shardMap.Version,
		Shards:     make(map[int][]string),
		Nodes:      make(map[string]NodeInfo),
		UpdateTime: sm.shardMap.UpdateTime,
	}

	// Copy the shard assignments
	for shardID, nodes := range sm.shardMap.Shards {
		copy.Shards[shardID] = append([]string(nil), nodes...)
	}

	// Copy the node info
	for nodeID, nodeInfo := range sm.shardMap.Nodes {
		copy.Nodes[nodeID] = nodeInfo
	}

	return copy
}

// RebalanceShards redistributes shards across available nodes
func (sm *ShardManager) RebalanceShards(nodes map[string]*NodeInfo) (*ShardMap, error) {
	if sm.placement == nil {
		return nil, fmt.Errorf("no placement strategy configured")
	}

	return sm.placement.RebalanceShards(sm.shardMap, nodes)
}

// PlaceActor determines which node should handle a new actor
func (sm *ShardManager) PlaceActor(actorID string, nodes map[string]*NodeInfo) (string, error) {
	if sm.placement == nil {
		return "", fmt.Errorf("no placement strategy configured")
	}

	return sm.placement.PlaceActor(actorID, sm.shardMap, nodes)
}

// GetActorsInShard returns actors that belong to a specific shard on a specific node
func (sm *ShardManager) GetActorsInShard(shardID int, nodeID string, vmRegistry VMRegistry) []string {
	if vmRegistry == nil {
		return []string{}
	}

	activeVMs := vmRegistry.GetActiveVMs()
	var actors []string

	for actorID := range activeVMs {
		if sm.GetShard(actorID) == shardID {
			actors = append(actors, actorID)
		}
	}

	return actors
}

// GetShardCount returns the total number of shards
func (sm *ShardManager) GetShardCount() int {
	return sm.shardCount
}

// GetReplicationFactor returns the replication factor
func (sm *ShardManager) GetReplicationFactor() int {
	return sm.replication
}

// ConsistentHashPlacement implements PlacementStrategy using consistent hashing
type ConsistentHashPlacement struct{}

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

	ring := NewConsistentHashRing()
	for nodeID := range nodes {
		ring.AddNode(nodeID)
	}

	node := ring.GetNode(actorID)
	if node == "" {
		sortedNodeIDs := make([]string, 0, len(nodes))
		for nodeID := range nodes {
			sortedNodeIDs = append(sortedNodeIDs, nodeID)
		}
		sort.Strings(sortedNodeIDs)
		return sortedNodeIDs[0], nil
	}

	return node, nil
}

// RebalanceShards redistributes shards across nodes using consistent hashing
func (chp *ConsistentHashPlacement) RebalanceShards(currentMap *ShardMap, nodes map[string]*NodeInfo) (*ShardMap, error) {
	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
}