seaweedfs/weed/mq/kafka/partition_mapping_test.go

package kafka

import (
	"testing"
	"time"

	"github.com/seaweedfs/seaweedfs/weed/mq/pub_balancer"
	"github.com/seaweedfs/seaweedfs/weed/pb/schema_pb"
)

// PartitionMapper provides consistent Kafka partition to SeaweedMQ ring mapping
// NOTE: This is test-only code and not used in the actual Kafka Gateway implementation
type PartitionMapper struct{}

// NewPartitionMapper creates a new partition mapper
func NewPartitionMapper() *PartitionMapper {
	return &PartitionMapper{}
}

// GetRangeSize returns the consistent range size for Kafka partition mapping
// This ensures all components use the same calculation
func (pm *PartitionMapper) GetRangeSize() int32 {
	// Use a range size that divides evenly into MaxPartitionCount (2520)
	// Range size 35 gives us exactly 72 Kafka partitions: 2520 / 35 = 72
	// This provides a good balance between partition granularity and ring utilization
	return 35
}

// GetMaxKafkaPartitions returns the maximum number of Kafka partitions supported
func (pm *PartitionMapper) GetMaxKafkaPartitions() int32 {
	// With range size 35, we can support: 2520 / 35 = 72 Kafka partitions
	return int32(pub_balancer.MaxPartitionCount) / pm.GetRangeSize()
}

// MapKafkaPartitionToSMQRange maps a Kafka partition to SeaweedMQ ring range
func (pm *PartitionMapper) MapKafkaPartitionToSMQRange(kafkaPartition int32) (rangeStart, rangeStop int32) {
	rangeSize := pm.GetRangeSize()
	rangeStart = kafkaPartition * rangeSize
	rangeStop = rangeStart + rangeSize - 1
	return rangeStart, rangeStop
}

// CreateSMQPartition creates a SeaweedMQ partition from a Kafka partition
func (pm *PartitionMapper) CreateSMQPartition(kafkaPartition int32, unixTimeNs int64) *schema_pb.Partition {
	rangeStart, rangeStop := pm.MapKafkaPartitionToSMQRange(kafkaPartition)

	return &schema_pb.Partition{
		RingSize:   pub_balancer.MaxPartitionCount,
		RangeStart: rangeStart,
		RangeStop:  rangeStop,
		UnixTimeNs: unixTimeNs,
	}
}

// ExtractKafkaPartitionFromSMQRange extracts the Kafka partition from SeaweedMQ range
func (pm *PartitionMapper) ExtractKafkaPartitionFromSMQRange(rangeStart int32) int32 {
	rangeSize := pm.GetRangeSize()
	return rangeStart / rangeSize
}

// ValidateKafkaPartition validates that a Kafka partition is within supported range
func (pm *PartitionMapper) ValidateKafkaPartition(kafkaPartition int32) bool {
	return kafkaPartition >= 0 && kafkaPartition < pm.GetMaxKafkaPartitions()
}

// GetPartitionMappingInfo returns debug information about the partition mapping
func (pm *PartitionMapper) GetPartitionMappingInfo() map[string]interface{} {
	return map[string]interface{}{
		"ring_size":            pub_balancer.MaxPartitionCount,
		"range_size":           pm.GetRangeSize(),
		"max_kafka_partitions": pm.GetMaxKafkaPartitions(),
		"ring_utilization":     float64(pm.GetMaxKafkaPartitions()*pm.GetRangeSize()) / float64(pub_balancer.MaxPartitionCount),
	}
}

// Global instance for consistent usage across the test codebase
var DefaultPartitionMapper = NewPartitionMapper()

func TestPartitionMapper_GetRangeSize(t *testing.T) {
	mapper := NewPartitionMapper()
	rangeSize := mapper.GetRangeSize()

	if rangeSize != 35 {
		t.Errorf("Expected range size 35, got %d", rangeSize)
	}

	// Verify that the range size divides evenly into available partitions
	maxPartitions := mapper.GetMaxKafkaPartitions()
	totalUsed := maxPartitions * rangeSize

	if totalUsed > int32(pub_balancer.MaxPartitionCount) {
		t.Errorf("Total used slots (%d) exceeds MaxPartitionCount (%d)", totalUsed, pub_balancer.MaxPartitionCount)
	}

	t.Logf("Range size: %d, Max Kafka partitions: %d, Ring utilization: %.2f%%",
		rangeSize, maxPartitions, float64(totalUsed)/float64(pub_balancer.MaxPartitionCount)*100)
}

func TestPartitionMapper_MapKafkaPartitionToSMQRange(t *testing.T) {
	mapper := NewPartitionMapper()

	tests := []struct {
		kafkaPartition int32
		expectedStart  int32
		expectedStop   int32
	}{
		{0, 0, 34},
		{1, 35, 69},
		{2, 70, 104},
		{10, 350, 384},
	}

	for _, tt := range tests {
		t.Run("", func(t *testing.T) {
			start, stop := mapper.MapKafkaPartitionToSMQRange(tt.kafkaPartition)

			if start != tt.expectedStart {
				t.Errorf("Kafka partition %d: expected start %d, got %d", tt.kafkaPartition, tt.expectedStart, start)
			}

			if stop != tt.expectedStop {
				t.Errorf("Kafka partition %d: expected stop %d, got %d", tt.kafkaPartition, tt.expectedStop, stop)
			}

			// Verify range size is consistent
			rangeSize := stop - start + 1
			if rangeSize != mapper.GetRangeSize() {
				t.Errorf("Inconsistent range size: expected %d, got %d", mapper.GetRangeSize(), rangeSize)
			}
		})
	}
}

func TestPartitionMapper_ExtractKafkaPartitionFromSMQRange(t *testing.T) {
	mapper := NewPartitionMapper()

	tests := []struct {
		rangeStart    int32
		expectedKafka int32
	}{
		{0, 0},
		{35, 1},
		{70, 2},
		{350, 10},
	}

	for _, tt := range tests {
		t.Run("", func(t *testing.T) {
			kafkaPartition := mapper.ExtractKafkaPartitionFromSMQRange(tt.rangeStart)

			if kafkaPartition != tt.expectedKafka {
				t.Errorf("Range start %d: expected Kafka partition %d, got %d",
					tt.rangeStart, tt.expectedKafka, kafkaPartition)
			}
		})
	}
}

func TestPartitionMapper_RoundTrip(t *testing.T) {
	mapper := NewPartitionMapper()

	// Test round-trip conversion for all valid Kafka partitions
	maxPartitions := mapper.GetMaxKafkaPartitions()

	for kafkaPartition := int32(0); kafkaPartition < maxPartitions; kafkaPartition++ {
		// Kafka -> SMQ -> Kafka
		rangeStart, rangeStop := mapper.MapKafkaPartitionToSMQRange(kafkaPartition)
		extractedKafka := mapper.ExtractKafkaPartitionFromSMQRange(rangeStart)

		if extractedKafka != kafkaPartition {
			t.Errorf("Round-trip failed for partition %d: got %d", kafkaPartition, extractedKafka)
		}

		// Verify no overlap with next partition
		if kafkaPartition < maxPartitions-1 {
			nextStart, _ := mapper.MapKafkaPartitionToSMQRange(kafkaPartition + 1)
			if rangeStop >= nextStart {
				t.Errorf("Partition %d range [%d,%d] overlaps with partition %d start %d",
					kafkaPartition, rangeStart, rangeStop, kafkaPartition+1, nextStart)
			}
		}
	}
}

func TestPartitionMapper_CreateSMQPartition(t *testing.T) {
	mapper := NewPartitionMapper()

	kafkaPartition := int32(5)
	unixTimeNs := time.Now().UnixNano()

	partition := mapper.CreateSMQPartition(kafkaPartition, unixTimeNs)

	if partition.RingSize != pub_balancer.MaxPartitionCount {
		t.Errorf("Expected ring size %d, got %d", pub_balancer.MaxPartitionCount, partition.RingSize)
	}

	expectedStart, expectedStop := mapper.MapKafkaPartitionToSMQRange(kafkaPartition)
	if partition.RangeStart != expectedStart {
		t.Errorf("Expected range start %d, got %d", expectedStart, partition.RangeStart)
	}

	if partition.RangeStop != expectedStop {
		t.Errorf("Expected range stop %d, got %d", expectedStop, partition.RangeStop)
	}

	if partition.UnixTimeNs != unixTimeNs {
		t.Errorf("Expected timestamp %d, got %d", unixTimeNs, partition.UnixTimeNs)
	}
}

func TestPartitionMapper_ValidateKafkaPartition(t *testing.T) {
	mapper := NewPartitionMapper()

	tests := []struct {
		partition int32
		valid     bool
	}{
		{-1, false},
		{0, true},
		{1, true},
		{mapper.GetMaxKafkaPartitions() - 1, true},
		{mapper.GetMaxKafkaPartitions(), false},
		{1000, false},
	}

	for _, tt := range tests {
		t.Run("", func(t *testing.T) {
			valid := mapper.ValidateKafkaPartition(tt.partition)
			if valid != tt.valid {
				t.Errorf("Partition %d: expected valid=%v, got %v", tt.partition, tt.valid, valid)
			}
		})
	}
}

func TestPartitionMapper_ConsistencyWithGlobalFunctions(t *testing.T) {
	mapper := NewPartitionMapper()

	kafkaPartition := int32(7)
	unixTimeNs := time.Now().UnixNano()

	// Test that global functions produce same results as mapper methods
	start1, stop1 := mapper.MapKafkaPartitionToSMQRange(kafkaPartition)
	start2, stop2 := MapKafkaPartitionToSMQRange(kafkaPartition)

	if start1 != start2 || stop1 != stop2 {
		t.Errorf("Global function inconsistent: mapper=(%d,%d), global=(%d,%d)",
			start1, stop1, start2, stop2)
	}

	partition1 := mapper.CreateSMQPartition(kafkaPartition, unixTimeNs)
	partition2 := CreateSMQPartition(kafkaPartition, unixTimeNs)

	if partition1.RangeStart != partition2.RangeStart || partition1.RangeStop != partition2.RangeStop {
		t.Errorf("Global CreateSMQPartition inconsistent")
	}

	extracted1 := mapper.ExtractKafkaPartitionFromSMQRange(start1)
	extracted2 := ExtractKafkaPartitionFromSMQRange(start1)

	if extracted1 != extracted2 {
		t.Errorf("Global ExtractKafkaPartitionFromSMQRange inconsistent: %d vs %d", extracted1, extracted2)
	}
}

func TestPartitionMapper_GetPartitionMappingInfo(t *testing.T) {
	mapper := NewPartitionMapper()

	info := mapper.GetPartitionMappingInfo()

	// Verify all expected keys are present
	expectedKeys := []string{"ring_size", "range_size", "max_kafka_partitions", "ring_utilization"}
	for _, key := range expectedKeys {
		if _, exists := info[key]; !exists {
			t.Errorf("Missing key in mapping info: %s", key)
		}
	}

	// Verify values are reasonable
	if info["ring_size"].(int) != pub_balancer.MaxPartitionCount {
		t.Errorf("Incorrect ring_size in info")
	}

	if info["range_size"].(int32) != mapper.GetRangeSize() {
		t.Errorf("Incorrect range_size in info")
	}

	utilization := info["ring_utilization"].(float64)
	if utilization <= 0 || utilization > 1 {
		t.Errorf("Invalid ring utilization: %f", utilization)
	}

	t.Logf("Partition mapping info: %+v", info)
}