From 8751292c2b7a5c5100a71ebec87f8590a0a77d0a Mon Sep 17 00:00:00 2001
From: Rajdeep Singh Chauhan <rajdchau@redhat.com>
Date: Tue, 20 Jun 2023 13:10:50 -0400
Subject: [PATCH] add vmss health check after vmss restart

---
 pkg/deploy/predeploy.go | 12 ++++++++++++
 1 file changed, 12 insertions(+)

diff --git a/pkg/deploy/predeploy.go b/pkg/deploy/predeploy.go
index 50ec43596b2..29aa89e7206 100644
--- a/pkg/deploy/predeploy.go
+++ b/pkg/deploy/predeploy.go
@@ -505,11 +505,14 @@ func (d *deployer) restartOldScalesets(ctx context.Context, resourceGroupName st
 
 func (d *deployer) restartOldScaleset(ctx context.Context, vmssName string, resourceGroupName string) error {
 	var restartScript string
+	var waitForReadiness func(ctx context.Context, vmssName string) error
 	switch {
 	case strings.HasPrefix(vmssName, gatewayVMSSPrefix):
 		restartScript = gatewayRestartScript
+		waitForReadiness = d.gatewayWaitForReadiness
 	case strings.HasPrefix(vmssName, rpVMSSPrefix):
 		restartScript = rpRestartScript
+		waitForReadiness = d.rpWaitForReadiness
 	default:
 		return nil
 	}
@@ -537,5 +540,14 @@ func (d *deployer) restartOldScaleset(ctx context.Context, vmssName string, reso
 			return err
 		}
 	}
+
+	// wait for load balancer probe to change the health status
+	time.Sleep(30 * time.Second)
+	timeoutCtx, cancel := context.WithTimeout(ctx, time.Hour)
+	defer cancel()
+	err = waitForReadiness(timeoutCtx, vmssName)
+	if err != nil {
+		return err
+	}
 	return nil
 }