From 7a95777731e7a07f5cbacaa56f419295071b0479 Mon Sep 17 00:00:00 2001
From: Jorres Tarasov <jorres.tarasov@gmail.com>
Date: Wed, 22 Jan 2025 12:21:45 +0100
Subject: [PATCH] wip

---
 pkg/rolling/restart_handler.go | 11 ++++-------
 pkg/rolling/rolling.go         | 11 ++++-------
 2 files changed, 8 insertions(+), 14 deletions(-)

diff --git a/pkg/rolling/restart_handler.go b/pkg/rolling/restart_handler.go
index 44da7ba..8ce3dee 100644
--- a/pkg/rolling/restart_handler.go
+++ b/pkg/rolling/restart_handler.go
@@ -89,14 +89,11 @@ func (rh *restartHandler) run() {
 	}
 }
 
-func (rh *restartHandler) stopWithDelay() {
-	close(rh.queue)
-	rh.wg.Wait()
-	close(rh.done)
-}
-
-func (rh *restartHandler) stopRightNow() {
+func (rh *restartHandler) stop(shouldWaitForDelay bool) {
 	close(rh.queue)
+	if shouldWaitForDelay {
+		rh.wg.Wait()
+	}
 	close(rh.done)
 }
 
diff --git a/pkg/rolling/rolling.go b/pkg/rolling/rolling.go
index 0f53e00..1f89a51 100644
--- a/pkg/rolling/rolling.go
+++ b/pkg/rolling/rolling.go
@@ -371,15 +371,12 @@ func (r *Rolling) processActionGroupStates(actions []*Ydb_Maintenance.ActionGrou
 	r.logCompleteResult(result)
 	r.state.unreportedButFinishedActionIds = []string{}
 
-	isRestartCompleted := len(actions) == len(result.ActionStatuses)
+	restartCompleted := len(actions) == len(result.ActionStatuses)
+	shouldWaitForDelay := !restartCompleted
 
-	if !isRestartCompleted {
-		restartHandler.stopWithDelay()
-	} else {
-		restartHandler.stopRightNow()
-	}
+	restartHandler.stop(shouldWaitForDelay)
 
-	return isRestartCompleted
+	return restartCompleted
 }
 
 func (r *Rolling) atomicHasActionInUnreported(actionID string) bool {