From cbfe0ab4f4c2a5ca8ee07da5dfc7fbc96978e27b Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Peter=20=C5=A0tibran=C3=BD?= <pstibrany@gmail.com>
Date: Fri, 16 Aug 2024 19:56:24 +0200
Subject: [PATCH] Fix computation of ShuffleShardExpectedInstancesPerZone for
 math.MaxInt. (#567)

---
 ring/shard/shard.go      | 3 +++
 ring/shard/shard_test.go | 6 ++++++
 2 files changed, 9 insertions(+)

diff --git a/ring/shard/shard.go b/ring/shard/shard.go
index 26d695514..bbc966707 100644
--- a/ring/shard/shard.go
+++ b/ring/shard/shard.go
@@ -30,6 +30,9 @@ func ShuffleShardSeed(identifier, zone string) int64 {
 // zone when zone-aware replication is enabled. The algorithm expects the shard size to be divisible
 // by the number of zones, in order to have nodes balanced across zones. If it's not, we do round up.
 func ShuffleShardExpectedInstancesPerZone(shardSize, numZones int) int {
+	if shardSize == math.MaxInt {
+		return math.MaxInt
+	}
 	return int(math.Ceil(float64(shardSize) / float64(numZones)))
 }
 
diff --git a/ring/shard/shard_test.go b/ring/shard/shard_test.go
index f23260bbc..63e84786b 100644
--- a/ring/shard/shard_test.go
+++ b/ring/shard/shard_test.go
@@ -1,6 +1,7 @@
 package shard
 
 import (
+	"math"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
@@ -38,6 +39,11 @@ func TestShuffleShardExpectedInstancesPerZone(t *testing.T) {
 			numZones:  3,
 			expected:  2,
 		},
+		{
+			shardSize: math.MaxInt,
+			numZones:  3,
+			expected:  math.MaxInt,
+		},
 	}
 
 	for _, test := range tests {