From 727b33c7106c16d8120592661032b0f3811b8abd Mon Sep 17 00:00:00 2001
From: Swagat Bora <sbora@amazon.com>
Date: Sat, 15 Oct 2022 00:54:28 +0000
Subject: [PATCH] Retry client connection in  waitForStart

Signed-off-by: Swagat Bora <sbora@amazon.com>
---
 integration/client/container_linux_test.go |  2 +-
 integration/client/daemon.go               | 32 ++++++++++++++--------
 2 files changed, 22 insertions(+), 12 deletions(-)

diff --git a/integration/client/container_linux_test.go b/integration/client/container_linux_test.go
index 6365ba740..be092d381 100644
--- a/integration/client/container_linux_test.go
+++ b/integration/client/container_linux_test.go
@@ -1006,7 +1006,7 @@ func TestDaemonRestartWithRunningShim(t *testing.T) {
 		t.Errorf(`first task.Wait() should have failed with "transport is closing"`)
 	}
 
-	waitCtx, cancel := context.WithTimeout(ctx, 1*time.Second)
+	waitCtx, cancel := context.WithTimeout(ctx, 4*time.Second)
 	c, err := ctrd.waitForStart(waitCtx)
 	cancel()
 	if err != nil {
diff --git a/integration/client/daemon.go b/integration/client/daemon.go
index 8b7e22fc7..769857792 100644
--- a/integration/client/daemon.go
+++ b/integration/client/daemon.go
@@ -24,6 +24,7 @@ import (
 	"runtime"
 	"sync"
 	"syscall"
+	"time"
 
 	. "github.com/containerd/containerd"
 	exec "golang.org/x/sys/execabs"
@@ -59,21 +60,30 @@ func (d *daemon) waitForStart(ctx context.Context) (*Client, error) {
 		client  *Client
 		serving bool
 		err     error
+		ticker  = time.NewTicker(500 * time.Millisecond)
 	)
+	defer ticker.Stop()
 
-	client, err = New(d.addr)
-	if err != nil {
-		return nil, err
-	}
-	serving, err = client.IsServing(ctx)
-	if !serving {
-		client.Close()
-		if err == nil {
-			err = errors.New("connection was successful but service is not available")
+	for {
+		select {
+		case <-ticker.C:
+			client, err = New(d.addr)
+			if err != nil {
+				continue
+			}
+			serving, err = client.IsServing(ctx)
+			if !serving {
+				client.Close()
+				if err == nil {
+					err = errors.New("connection was successful but service is not available")
+				}
+				continue
+			}
+			return client, err
+		case <-ctx.Done():
+			return nil, fmt.Errorf("context deadline exceeded: %w", err)
 		}
-		return nil, err
 	}
-	return client, err
 }
 
 func (d *daemon) Stop() error {