docker-container: give readiness wait a full startup timeout

driver/docker-container/driver.go

@@ -35,9 +35,10 @@ import (
 )
 
 const (
-	volumeStateSuffix       = "_state"
-	buildkitdConfigFile     = "buildkitd.toml"
-	buildkitdStartupTimeout = 20 * time.Second
+	volumeStateSuffix      = "_state"
+	buildkitdConfigFile    = "buildkitd.toml"
+	buildkitdStartupWindow = 20 * time.Second
+	buildkitdReadyTimeout  = 20 * time.Second
 )
 
 type Driver struct {
@@ -525,7 +526,7 @@ func (d *Driver) Client(ctx context.Context, opts ...client.ClientOpt) (*client.
 		}
 		return nil, errors.WithStack(err)
 	}
-	waitDeadline, err := clientWaitDeadline(res.Container.State, time.Now())
+	waitReady, err := clientWaitReady(res.Container.State, time.Now())
 	if err != nil {
 		return nil, err
 	}
@@ -549,11 +550,11 @@ func (d *Driver) Client(ctx context.Context, opts ...client.ClientOpt) (*client.
 		_ = conn.Close()
 		return nil, err
 	}
-	if waitDeadline.IsZero() {
+	if !waitReady {
 		return c, nil
 	}
 
-	waitCtx, cancel := context.WithDeadlineCause(ctx, waitDeadline, errors.WithStack(context.DeadlineExceeded))
+	waitCtx, cancel := context.WithTimeoutCause(ctx, buildkitdReadyTimeout, errors.WithStack(context.DeadlineExceeded))
 	defer cancel()
 	if err := c.Wait(waitCtx); err != nil {
 		_ = c.Close()
@@ -562,22 +563,18 @@ func (d *Driver) Client(ctx context.Context, opts ...client.ClientOpt) (*client.
 	return c, nil
 }
 
-func clientWaitDeadline(state *container.State, now time.Time) (time.Time, error) {
+func clientWaitReady(state *container.State, now time.Time) (bool, error) {
 	if state == nil || !state.Running {
-		return time.Time{}, driver.ErrNotRunning{}
+		return false, 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.
+	// socket. Use the startup window only to decide whether to wait 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
+		return false, nil
 	}
-	deadline := startedAt.Add(buildkitdStartupTimeout)
-	if !now.Before(deadline) {
-		return time.Time{}, nil
-	}
-	return deadline, nil
+	return now.Before(startedAt.Add(buildkitdStartupWindow)), nil
 }
 
 func (d *Driver) Factory() driver.Factory {

driver/docker-container/driver_test.go

@@ -9,40 +9,35 @@ import (
 	"github.com/stretchr/testify/require"
 )
 
-func TestClientWaitDeadline(t *testing.T) {
+func TestClientWaitReady(t *testing.T) {
 	now := time.Now()
+	running := func(startedAt string) *container.State {
+		return &container.State{Running: true, StartedAt: startedAt}
+	}
+	tests := []struct {
+		name     string
+		state    *container.State
+		wantWait bool
+		wantErr  error
+	}{
+		{name: "missing-state-fails-fast", wantErr: driver.ErrNotRunning{}},
+		{name: "stopped-builder-fails-fast", state: &container.State{}, wantErr: driver.ErrNotRunning{}},
+		{name: "established-builder-skips-wait", state: running(now.Add(-2 * buildkitdStartupWindow).Format(time.RFC3339Nano))},
+		{name: "recent-start-builder-waits", state: running(now.Add(-time.Second).Format(time.RFC3339Nano)), wantWait: true},
+		{name: "nearly-expired-startup-window-waits", state: running(now.Add(-buildkitdStartupWindow + time.Nanosecond).Format(time.RFC3339Nano)), wantWait: true},
+		{name: "expired-startup-window-skips-wait", state: running(now.Add(-buildkitdStartupWindow).Format(time.RFC3339Nano))},
+		{name: "invalid-start-time-skips-wait", state: running("invalid")},
+	}
 
-	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())
-	})
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			wait, err := clientWaitReady(tt.state, now)
+			if tt.wantErr != nil {
+				require.ErrorIs(t, err, tt.wantErr)
+			} else {
+				require.NoError(t, err)
+			}
+			require.Equal(t, tt.wantWait, wait)
+		})
+	}
 }