docker-container: wait for BuildKit client readiness

driver/docker-container/driver.go

@@ -35,8 +35,9 @@ import (
 )
 
 const (
-	volumeStateSuffix   = "_state"
-	buildkitdConfigFile = "buildkitd.toml"
+	volumeStateSuffix       = "_state"
+	buildkitdConfigFile     = "buildkitd.toml"
+	buildkitdStartupTimeout = 20 * time.Second
 )
 
 type Driver struct {
@@ -517,21 +518,66 @@ func (d *Driver) Dial(ctx context.Context) (net.Conn, error) {
 }
 
 func (d *Driver) Client(ctx context.Context, opts ...client.ClientOpt) (*client.Client, error) {
+	res, err := d.DockerAPI.ContainerInspect(ctx, d.Name, dockerclient.ContainerInspectOptions{})
+	if err != nil {
+		if cerrdefs.IsNotFound(err) {
+			return nil, driver.ErrNotRunning{}
+		}
+		return nil, errors.WithStack(err)
+	}
+	waitDeadline, err := clientWaitDeadline(res.Container.State, time.Now())
+	if err != nil {
+		return nil, err
+	}
+
 	conn, err := d.Dial(ctx)
 	if err != nil {
 		return nil, err
 	}
 
 	var counter int64
 	opts = append([]client.ClientOpt{
-		client.WithContextDialer(func(context.Context, string) (net.Conn, error) {
-			if atomic.AddInt64(&counter, 1) > 1 {
-				return nil, net.ErrClosed
+		client.WithContextDialer(func(ctx context.Context, _ string) (net.Conn, error) {
+			if atomic.AddInt64(&counter, 1) == 1 {
+				return conn, nil
 			}
-			return conn, nil
+			return d.Dial(ctx)
 		}),
 	}, opts...)
-	return client.New(ctx, "", opts...)
+	c, err := client.New(ctx, "", opts...)
+	if err != nil {
+		_ = conn.Close()
+		return nil, err
+	}
+	if waitDeadline.IsZero() {
+		return c, nil
+	}
+
+	waitCtx, cancel := context.WithDeadlineCause(ctx, waitDeadline, errors.WithStack(context.DeadlineExceeded))
+	defer cancel()
+	if err := c.Wait(waitCtx); err != nil {
+		_ = c.Close()
+		return nil, errors.Wrap(err, "waiting for BuildKit")
+	}
+	return c, nil
+}
+
+func clientWaitDeadline(state *container.State, now time.Time) (time.Time, error) {
+	if state == nil || !state.Running {
+		return time.Time{}, driver.ErrNotRunning{}
+	}
+	// Docker reports a container as running before buildkitd has bound its
+	// socket. Wait only during that startup window so an established but broken
+	// builder still returns its connection error promptly.
+	startedAt, err := time.Parse(time.RFC3339Nano, state.StartedAt)
+	if err != nil {
+		return time.Time{}, nil
+	}
+	deadline := startedAt.Add(buildkitdStartupTimeout)
+	if !now.Before(deadline) {
+		return time.Time{}, nil
+	}
+	return deadline, nil
 }
 
 func (d *Driver) Factory() driver.Factory {

driver/docker-container/driver_test.go

@@ -0,0 +1,48 @@
+package docker
+
+import (
+	"testing"
+	"time"
+
+	"github.com/docker/buildx/driver"
+	"github.com/moby/moby/api/types/container"
+	"github.com/stretchr/testify/require"
+)
+
+func TestClientWaitDeadline(t *testing.T) {
+	now := time.Now()
+
+	t.Run("stopped-builder-fails-fast", func(t *testing.T) {
+		deadline, err := clientWaitDeadline(&container.State{}, now)
+		require.ErrorIs(t, err, driver.ErrNotRunning{})
+		require.True(t, deadline.IsZero())
+	})
+
+	t.Run("established-builder-skips-wait", func(t *testing.T) {
+		deadline, err := clientWaitDeadline(&container.State{
+			Running:   true,
+			StartedAt: now.Add(-2 * buildkitdStartupTimeout).Format(time.RFC3339Nano),
+		}, now)
+		require.NoError(t, err)
+		require.True(t, deadline.IsZero())
+	})
+
+	t.Run("recent-start-builder-waits", func(t *testing.T) {
+		startedAt := now.Add(-time.Second)
+		deadline, err := clientWaitDeadline(&container.State{
+			Running:   true,
+			StartedAt: startedAt.Format(time.RFC3339Nano),
+		}, now)
+		require.NoError(t, err)
+		require.True(t, deadline.Equal(startedAt.Add(buildkitdStartupTimeout)))
+	})
+
+	t.Run("invalid-start-time-skips-wait", func(t *testing.T) {
+		deadline, err := clientWaitDeadline(&container.State{
+			Running:   true,
+			StartedAt: "invalid",
+		}, now)
+		require.NoError(t, err)
+		require.True(t, deadline.IsZero())
+	})
+}

driver/manager.go

@@ -127,17 +127,23 @@ func GetFactories(instanceRequired bool) []Factory {
 type DriverHandle struct {
 	Driver
 	client                  *client.Client
-	err                     error
-	once                    sync.Once
+	clientMu                sync.Mutex
 	historyAPISupportedOnce sync.Once
 	historyAPISupported     bool
 }
 
 func (d *DriverHandle) Client(ctx context.Context, opt ...client.ClientOpt) (*client.Client, error) {
-	d.once.Do(func() {
-		d.client, d.err = d.Driver.Client(ctx, append(d.getClientOptions(), opt...)...)
-	})
-	return d.client, d.err
+	d.clientMu.Lock()
+	defer d.clientMu.Unlock()
+	if d.client != nil {
+		return d.client, nil
+	}
+	c, err := d.Driver.Client(ctx, append(d.getClientOptions(), opt...)...)
+	if err != nil {
+		return nil, err
+	}
+	d.client = c
+	return c, nil
 }
 
 func (d *DriverHandle) UncachedClient(ctx context.Context) (*client.Client, error) {

driver/manager_test.go

@@ -0,0 +1,36 @@
+package driver
+
+import (
+	"context"
+	"testing"
+
+	"github.com/moby/buildkit/client"
+	"github.com/stretchr/testify/require"
+)
+
+func TestBootRetriesClientAfterErrNotRunning(t *testing.T) {
+	d := &retryDriver{client: &client.Client{}}
+
+	c, err := Boot(context.Background(), context.Background(), &DriverHandle{Driver: d}, nil)
+	require.NoError(t, err)
+	require.Same(t, d.client, c)
+	require.Equal(t, 2, d.clientCalls)
+}
+
+type retryDriver struct {
+	Driver
+	client      *client.Client
+	clientCalls int
+}
+
+func (d *retryDriver) Info(context.Context) (*Info, error) {
+	return &Info{Status: Running}, nil
+}
+
+func (d *retryDriver) Client(context.Context, ...client.ClientOpt) (*client.Client, error) {
+	d.clientCalls++
+	if d.clientCalls == 1 {
+		return nil, ErrNotRunning{}
+	}
+	return d.client, nil
+}