From 597606288b570313420807a19ce44dbfc0532332 Mon Sep 17 00:00:00 2001
From: Justin Chadwell <me@jedevc.com>
Date: Tue, 5 Sep 2023 13:10:30 +0100
Subject: [PATCH] client: manually implement Wait backoffs

When calling client.Wait, we want to avoid the default backoff behavior,
because we want to achieve a quick response back once the server becomes
active.

To do this, without modifying the entire client's exponential backoff
configuration, we can use conn.ResetConnectBackoff, while attempting to
reconnect every second.

Here are some common scenarios:
- Server is listening: the call to Info succeeds quickly, and we return.
- Server is listening, but is behind several proxies and so latency is
  high: the call to Info succeeds slowly (up to minConnectTimeout=20s),
  and we return.
- Server is not listening and gets "connection refused": the
  call to Info fails quickly, and we wait a second before retrying.
- Server is not listening and does not respond (e.g. firewall dropping
  packets): the call to Info fails slowly (by default after
  minConnectTimeout=20s). After the call fails, we wait a second before
  retrying.

Signed-off-by: Justin Chadwell <me@jedevc.com>
(cherry picked from commit f1d7f2e410c78922f42304830cf2164238a651e8)
Signed-off-by: Justin Chadwell <me@jedevc.com>
---
 client/client.go | 24 +++++++++++++++++++-----
 1 file changed, 19 insertions(+), 5 deletions(-)

diff --git a/client/client.go b/client/client.go
index 1d60a7068340..71a72bf9f6a5 100644
--- a/client/client.go
+++ b/client/client.go
@@ -8,6 +8,7 @@ import (
 	"net/url"
 	"os"
 	"strings"
+	"time"
 
 	contentapi "github.com/containerd/containerd/api/services/content/v1"
 	"github.com/containerd/containerd/defaults"
@@ -186,16 +187,29 @@ func (c *Client) Dialer() session.Dialer {
 }
 
 func (c *Client) Wait(ctx context.Context) error {
-	opts := []grpc.CallOption{grpc.WaitForReady(true)}
-	_, err := c.ControlClient().Info(ctx, &controlapi.InfoRequest{}, opts...)
-	if err != nil {
-		if code := grpcerrors.Code(err); code == codes.Unimplemented {
+	for {
+		_, err := c.ControlClient().Info(ctx, &controlapi.InfoRequest{})
+		if err == nil {
+			return nil
+		}
+
+		switch code := grpcerrors.Code(err); code {
+		case codes.Unavailable:
+		case codes.Unimplemented:
 			// only buildkit v0.11+ supports the info api, but an unimplemented
 			// response error is still a response so we can ignore it
 			return nil
+		default:
+			return err
+		}
+
+		select {
+		case <-ctx.Done():
+			return ctx.Err()
+		case <-time.After(time.Second):
 		}
+		c.conn.ResetConnectBackoff()
 	}
-	return err
 }
 
 func (c *Client) Close() error {