address partition overflow bug

flow/connectors/mongo/qrep_partition.go

@@ -139,7 +139,7 @@ func (c *MongoConnector) numericPartitions(
 		return utils.FullTablePartition(), nil
 	}
 
-	ranges := computeNumericRanges(minVal, maxVal, numPartitions)
+	ranges := utils.ComputeRanges(minVal, maxVal, numPartitions)
 	c.logger.Info("[mongo] using numeric _id partitioning", slog.Int("numPartitions", len(ranges)))
 	partitions := make([]*protos.QRepPartition, 0, len(ranges))
 	for i, rng := range ranges {
@@ -148,27 +148,6 @@ func (c *MongoConnector) numericPartitions(
 	return partitions, nil
 }
 
-func computeNumericRanges(minVal int64, maxVal int64, numPartitions int64) [][2]int64 {
-	// span arithmetic is done in uint64 to avoid overflow
-	span := uint64(maxVal) - uint64(minVal)
-	step := span / uint64(numPartitions)
-	if span%uint64(numPartitions) != 0 {
-		step++
-	}
-
-	ranges := make([][2]int64, 0, numPartitions)
-	start := minVal
-	for {
-		if remaining := uint64(maxVal) - uint64(start); remaining <= step {
-			ranges = append(ranges, [2]int64{start, maxVal})
-			return ranges
-		}
-		end := int64(uint64(start) + step)
-		ranges = append(ranges, [2]int64{start, end})
-		start = end
-	}
-}
-
 // stringPartitions builds partitions for a string _id collection. Because string
 // keys (e.g. package names) are not uniformly distributed, uniform division of the
 // keyspace would be badly skewed; instead we sample the collection and pick quantile

flow/connectors/mongo/qrep_partition_test.go

@@ -1,7 +1,6 @@
 package connmongo
 
 import (
-	"math"
 	"testing"
 
 	"github.com/stretchr/testify/require"
@@ -119,51 +118,6 @@ func TestComputeStringBoundaries(t *testing.T) {
 	})
 }
 
-func assertContiguousNumeric(t *testing.T, ranges [][2]int64, minVal, maxVal int64) {
-	t.Helper()
-	require.NotEmpty(t, ranges)
-	require.Equal(t, minVal, ranges[0][0])
-	require.Equal(t, maxVal, ranges[len(ranges)-1][1])
-	for i := range ranges {
-		require.Less(t, ranges[i][0], ranges[i][1])
-		if i+1 < len(ranges) {
-			require.Equal(t, ranges[i][1], ranges[i+1][0])
-		}
-	}
-}
-
-func TestComputeNumericRanges(t *testing.T) {
-	t.Run("even division", func(t *testing.T) {
-		ranges := computeNumericRanges(1, 100, 10)
-		require.Len(t, ranges, 10)
-		assertContiguousNumeric(t, ranges, 1, 100)
-	})
-
-	t.Run("span smaller than partition count", func(t *testing.T) {
-		ranges := computeNumericRanges(0, 5, 10)
-		require.Len(t, ranges, 5)
-		assertContiguousNumeric(t, ranges, 0, 5)
-	})
-
-	t.Run("two adjacent values", func(t *testing.T) {
-		ranges := computeNumericRanges(7, 8, 4)
-		require.Len(t, ranges, 1)
-		assertContiguousNumeric(t, ranges, 7, 8)
-	})
-
-	t.Run("span does not overflow", func(t *testing.T) {
-		ranges := computeNumericRanges(math.MinInt64+5, math.MaxInt64-5, 10)
-		require.Len(t, ranges, 10)
-		assertContiguousNumeric(t, ranges, math.MinInt64+5, math.MaxInt64-5)
-	})
-
-	t.Run("full int64 domain does not overflow", func(t *testing.T) {
-		ranges := computeNumericRanges(math.MinInt64, math.MaxInt64, 7)
-		require.Len(t, ranges, 7)
-		assertContiguousNumeric(t, ranges, math.MinInt64, math.MaxInt64)
-	})
-}
-
 func TestToRangeFilter(t *testing.T) {
 	t.Run("numeric range is half-open", func(t *testing.T) {
 		filter, err := toRangeFilter(DefaultDocumentKeyColumnName, utils.CreateNumericPartition(10, 20, false).Range)

flow/connectors/utils/partition.go

@@ -12,7 +12,6 @@ import (
 	"google.golang.org/protobuf/types/known/timestamppb"
 
 	"github.com/PeerDB-io/peerdb/flow/generated/protos"
-	"github.com/PeerDB-io/peerdb/flow/shared"
 )
 
 const FullTablePartitionID = "full-table-partition-id"
@@ -314,32 +313,50 @@ func (p *PartitionHelper) AddPartitionsWithRange(start any, end any, numPartitio
 
 	switch r := partition.Range.Range.(type) {
 	case *protos.PartitionRange_IntRange:
-		size := shared.DivCeil(r.IntRange.End-r.IntRange.Start, numPartitions)
-		for i := range numPartitions {
-			if err := p.AddPartition(r.IntRange.Start+size*i, min(r.IntRange.Start+size*(i+1), r.IntRange.End)); err != nil {
+		for _, rng := range ComputeRanges(r.IntRange.Start, r.IntRange.End, numPartitions) {
+			if err := p.AddPartition(rng[0], rng[1]); err != nil {
 				return err
 			}
 		}
 	case *protos.PartitionRange_UintRange:
-		size := shared.DivCeil(r.UintRange.End-r.UintRange.Start, uint64(numPartitions))
-		for i := range uint64(numPartitions) {
-			if err := p.AddPartition(r.UintRange.Start+size*i, min(r.UintRange.Start+size*(i+1), r.UintRange.End)); err != nil {
+		for _, rng := range ComputeRanges(r.UintRange.Start, r.UintRange.End, numPartitions) {
+			if err := p.AddPartition(rng[0], rng[1]); err != nil {
 				return err
 			}
 		}
 	case *protos.PartitionRange_TimestampRange:
 		tstart := r.TimestampRange.Start.AsTime().UnixMicro()
 		tend := r.TimestampRange.End.AsTime().UnixMicro()
-		size := shared.DivCeil(tend-tstart, numPartitions)
-		for i := range numPartitions {
-			if err := p.AddPartition(time.UnixMicro(tstart+size*i), time.UnixMicro(min(tstart+size*(i+1), tend))); err != nil {
+		for _, rng := range ComputeRanges(tstart, tend, numPartitions) {
+			if err := p.AddPartition(time.UnixMicro(rng[0]), time.UnixMicro(rng[1])); err != nil {
 				return err
 			}
 		}
 	}
 	return nil
 }
 
+func ComputeRanges[T int64 | uint64](minVal T, maxVal T, numPartitions int64) [][2]T {
+	// span arithmetic is done in uint64 to avoid overflow
+	span := uint64(maxVal) - uint64(minVal)
+	step := span / uint64(numPartitions)
+	if span%uint64(numPartitions) != 0 {
+		step++
+	}
+
+	ranges := make([][2]T, 0, numPartitions)
+	start := minVal
+	for {
+		if remaining := uint64(maxVal) - uint64(start); remaining <= step {
+			ranges = append(ranges, [2]T{start, maxVal})
+			return ranges
+		}
+		end := T(uint64(start) + step)
+		ranges = append(ranges, [2]T{start, end})
+		start = end
+	}
+}
+
 func (p *PartitionHelper) getPartitionForStartAndEnd(start any, end any) (*protos.QRepPartition, error) {
 	if start == nil || end == nil {
 		return nil, nil

flow/connectors/utils/partition_test.go

@@ -0,0 +1,138 @@
+package utils
+
+import (
+	"log/slog"
+	"math"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+	"go.temporal.io/sdk/log"
+
+	"github.com/PeerDB-io/peerdb/flow/generated/protos"
+)
+
+func intRangeOf(t *testing.T, partition *protos.QRepPartition) *protos.IntPartitionRange {
+	t.Helper()
+	r, ok := partition.Range.Range.(*protos.PartitionRange_IntRange)
+	require.True(t, ok)
+	return r.IntRange
+}
+
+func uintRangeOf(t *testing.T, partition *protos.QRepPartition) *protos.UIntPartitionRange {
+	t.Helper()
+	r, ok := partition.Range.Range.(*protos.PartitionRange_UintRange)
+	require.True(t, ok)
+	return r.UintRange
+}
+
+func assertContiguousRanges[T int64 | uint64](t *testing.T, ranges [][2]T, minVal, maxVal T) {
+	t.Helper()
+	require.NotEmpty(t, ranges)
+	require.Equal(t, minVal, ranges[0][0])
+	require.Equal(t, maxVal, ranges[len(ranges)-1][1])
+	for i := range ranges {
+		require.Less(t, ranges[i][0], ranges[i][1])
+		if i+1 < len(ranges) {
+			require.Equal(t, ranges[i][1], ranges[i+1][0])
+		}
+	}
+}
+
+func TestComputeRanges(t *testing.T) {
+	t.Run("even division", func(t *testing.T) {
+		ranges := ComputeRanges[int64](1, 100, 10)
+		require.Len(t, ranges, 10)
+		assertContiguousRanges(t, ranges, int64(1), int64(100))
+	})
+
+	t.Run("span smaller than partition count", func(t *testing.T) {
+		ranges := ComputeRanges[int64](0, 5, 10)
+		require.Len(t, ranges, 5)
+		assertContiguousRanges(t, ranges, int64(0), int64(5))
+	})
+
+	t.Run("two adjacent values", func(t *testing.T) {
+		ranges := ComputeRanges[int64](7, 8, 4)
+		require.Len(t, ranges, 1)
+		assertContiguousRanges(t, ranges, int64(7), int64(8))
+	})
+
+	t.Run("full int64 domain does not overflow", func(t *testing.T) {
+		ranges := ComputeRanges[int64](math.MinInt64, math.MaxInt64, 7)
+		require.Len(t, ranges, 7)
+		assertContiguousRanges(t, ranges, int64(math.MinInt64), int64(math.MaxInt64))
+	})
+
+	t.Run("full uint64 domain does not overflow", func(t *testing.T) {
+		ranges := ComputeRanges[uint64](0, math.MaxUint64, 7)
+		require.Len(t, ranges, 7)
+		assertContiguousRanges(t, ranges, uint64(0), uint64(math.MaxUint64))
+	})
+
+	t.Run("same min/max values", func(t *testing.T) {
+		ranges := ComputeRanges[int64](7, 7, 4)
+		require.Len(t, ranges, 1)
+		require.Equal(t, int64(7), ranges[0][0])
+		require.Equal(t, int64(7), ranges[len(ranges)-1][1])
+	})
+}
+
+func TestAddPartitionsWithRange(t *testing.T) {
+	logger := log.NewStructuredLogger(slog.Default())
+
+	t.Run("negative to positive int range covers domain contiguously", func(t *testing.T) {
+		helper := NewPartitionHelper(logger)
+		require.NoError(t, helper.AddPartitionsWithRange(int64(-100), int64(100), 4))
+
+		partitions := helper.GetPartitions()
+		require.Len(t, partitions, 4)
+		require.Equal(t, int64(-100), intRangeOf(t, partitions[0]).Start)
+		require.Equal(t, int64(100), intRangeOf(t, partitions[len(partitions)-1]).End)
+		for i := 1; i < len(partitions); i++ {
+			require.Equal(t, intRangeOf(t, partitions[i-1]).End+1, intRangeOf(t, partitions[i]).Start)
+		}
+	})
+
+	t.Run("int range spanning full int64 domain does not overflow", func(t *testing.T) {
+		helper := NewPartitionHelper(logger)
+		require.NoError(t, helper.AddPartitionsWithRange(int64(math.MinInt64), int64(math.MaxInt64), 7))
+
+		partitions := helper.GetPartitions()
+		require.Len(t, partitions, 7)
+		require.Equal(t, int64(math.MinInt64), intRangeOf(t, partitions[0]).Start)
+		require.Equal(t, int64(math.MaxInt64), intRangeOf(t, partitions[len(partitions)-1]).End)
+		for i, partition := range partitions {
+			r := intRangeOf(t, partition)
+			require.LessOrEqual(t, r.Start, r.End)
+			if i > 0 {
+				require.Equal(t, intRangeOf(t, partitions[i-1]).End+1, r.Start)
+			}
+		}
+	})
+
+	t.Run("uint range spanning full uint64 domain does not overflow", func(t *testing.T) {
+		helper := NewPartitionHelper(logger)
+		require.NoError(t, helper.AddPartitionsWithRange(uint64(0), uint64(math.MaxUint64), 7))
+
+		partitions := helper.GetPartitions()
+		require.Len(t, partitions, 7)
+		require.Equal(t, uint64(0), uintRangeOf(t, partitions[0]).Start)
+		require.Equal(t, uint64(math.MaxUint64), uintRangeOf(t, partitions[len(partitions)-1]).End)
+		for i, partition := range partitions {
+			r := uintRangeOf(t, partition)
+			require.LessOrEqual(t, r.Start, r.End)
+			if i > 0 {
+				require.Equal(t, uintRangeOf(t, partitions[i-1]).End+1, r.Start)
+			}
+		}
+	})
+
+	t.Run("span smaller than partition count yields fewer partitions", func(t *testing.T) {
+		helper := NewPartitionHelper(logger)
+		require.NoError(t, helper.AddPartitionsWithRange(int64(0), int64(5), 10))
+
+		partitions := helper.GetPartitions()
+		require.Len(t, partitions, 5)
+		require.Equal(t, int64(5), intRangeOf(t, partitions[len(partitions)-1]).End)
+	})
+}