Retry client connection in waitForStart

Signed-off-by: Swagat Bora <sbora@amazon.com>

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 {

integration/client/daemon.go

@@ -24,6 +24,7 @@ import (
 	"runtime"
 	"sync"
 	"syscall"
+	"time"
 
 	. "github.com/containerd/containerd"
 	exec "golang.org/x/sys/execabs"
@@ -59,11 +60,16 @@ func (d *daemon) waitForStart(ctx context.Context) (*Client, error) {
 		client  *Client
 		serving bool
 		err     error
+		ticker  = time.NewTicker(500 * time.Millisecond)
 	)
+	defer ticker.Stop()
 
+	for {
+		select {
+		case <-ticker.C:
 			client, err = New(d.addr)
 			if err != nil {
-		return nil, err
+				continue
 			}
 			serving, err = client.IsServing(ctx)
 			if !serving {
@@ -71,9 +77,13 @@ func (d *daemon) waitForStart(ctx context.Context) (*Client, error) {
 				if err == nil {
 					err = errors.New("connection was successful but service is not available")
 				}
-		return nil, err
+				continue
 			}
 			return client, err
+		case <-ctx.Done():
+			return nil, fmt.Errorf("context deadline exceeded: %w", err)
+		}
+	}
 }
 
 func (d *daemon) Stop() error {