From f8de577747a9fd65a45a97e7b0e68d58ac2b80d1 Mon Sep 17 00:00:00 2001
From: Christopher Li <christopher.li@databricks.com>
Date: Thu, 16 May 2024 15:17:22 -0700
Subject: [PATCH] waited for all pods ready during generated hashring
 initialization (#3)

* return when encountering error

* waited for all pods ready during generated hashring initialization
---
 .golangci.yml |  1 +
 main.go       | 20 +++++++++++++++++---
 2 files changed, 18 insertions(+), 3 deletions(-)

diff --git a/.golangci.yml b/.golangci.yml
index 5b5b652..7985367 100644
--- a/.golangci.yml
+++ b/.golangci.yml
@@ -24,6 +24,7 @@ linters:
   - perfsprint
   - maligned
   - gosec
+  - gocognit
 
 linters-settings:
   errcheck:
diff --git a/main.go b/main.go
index 0da18ea..bae0880 100644
--- a/main.go
+++ b/main.go
@@ -571,10 +571,24 @@ func (c *controller) sync(ctx context.Context) {
 			continue
 		}
 
-		// If there's an increase in replicas we poll for the new replicas to be ready
-		if _, ok := c.replicas[sts.Name]; ok && c.replicas[sts.Name] < *sts.Spec.Replicas {
+		stsReplica, exist := c.replicas[sts.Name]
+		// If hashring is not initialized, need to wait for all pods ready within statefulset before generating hashring
+		if !exist && c.options.allowOnlyReadyReplicas {
+			for i := int32(0); i < *sts.Spec.Replicas; i++ {
+				start := time.Now()
+				podName := fmt.Sprintf("%s-%d", sts.Name, i)
+
+				if err := c.waitForPod(ctx, podName); err != nil {
+					level.Warn(c.logger).Log("msg", "failed waiting for pod ready during hashring intialization", "pod", podName, "duration", time.Since(start), "err", err)
+					return
+				}
+
+				level.Debug(c.logger).Log("msg", "waited until new pod was ready during hashring intialization", "pod", podName, "duration", time.Since(start))
+			}
+		} else if exist && stsReplica < *sts.Spec.Replicas {
+			// If there's an increase in replicas we poll for the new replicas to be ready
 			// Iterate over new replicas to wait until they are running
-			for i := c.replicas[sts.Name]; i < *sts.Spec.Replicas; i++ {
+			for i := stsReplica; i < *sts.Spec.Replicas; i++ {
 				start := time.Now()
 				podName := fmt.Sprintf("%s-%d", sts.Name, i)