aether/cluster/shard.go

package cluster

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

// 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
func NewShardManager(shardCount, replication int) *ShardManager {
	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
}


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

// PlaceActor places an actor using consistent hashing
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
	for nodeID := range nodes {
		if i == int(nodeIndex) {
			return nodeID, nil
		}
		i++
	}

	// Fallback to first node
	for nodeID := range nodes {
		return nodeID, nil
	}

	return "", fmt.Errorf("failed to place actor")
}

// RebalanceShards rebalances shards across nodes
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
}