add timeout parameter to cli and driver

avoidik · avoidik · commit c33a03a68a9b · 2024-07-10T15:28:54.000+01:00
Signed-off-by: avoidik &lt;avoidik@gmail.com&gt;
diff --git a/builder/builder.go b/builder/builder.go
@@ -343,6 +343,7 @@ type CreateOpts struct {
 	Use                 bool
 	Endpoint            string
 	Append              bool
+	Timeout             time.Duration
 }
 
 func Create(ctx context.Context, txn *store.Txn, dockerCli command.Cli, opts CreateOpts) (*Builder, error) {
@@ -522,7 +523,7 @@ func Create(ctx context.Context, txn *store.Txn, dockerCli command.Cli, opts Cre
 		return nil, err
 	}
 
-	timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
+	timeoutCtx, cancel := context.WithTimeout(ctx, opts.Timeout)
 	defer cancel()
 
 	nodes, err := b.LoadNodes(timeoutCtx, WithData())
diff --git a/commands/build.go b/commands/build.go
@@ -104,6 +104,8 @@ type buildOptions struct {
 	control.ControlOptions
 
 	invokeConfig *invokeConfig
+
+	timeout time.Duration
 }
 
 func (o *buildOptions) toControllerOptions() (*controllerapi.BuildOptions, error) {
diff --git a/commands/create.go b/commands/create.go
@@ -4,6 +4,7 @@ import (
 	"bytes"
 	"context"
 	"fmt"
+	"time"
 
 	"github.com/docker/buildx/builder"
 	"github.com/docker/buildx/driver"
@@ -27,6 +28,7 @@ type createOptions struct {
 	buildkitdFlags      string
 	buildkitdConfigFile string
 	bootstrap           bool
+	timeout             time.Duration
 	// upgrade      bool // perform upgrade of the driver
 }
 
@@ -61,6 +63,7 @@ func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, arg
 		Use:                 in.use,
 		Endpoint:            ep,
 		Append:              in.actionAppend,
+		Timeout:             in.timeout,
 	})
 	if err != nil {
 		return err
@@ -80,7 +83,7 @@ func runCreate(ctx context.Context, dockerCli command.Cli, in createOptions, arg
 	return nil
 }
 
-func createCmd(dockerCli command.Cli) *cobra.Command {
+func createCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 	var options createOptions
 
 	var drivers bytes.Buffer
@@ -96,6 +99,7 @@ func createCmd(dockerCli command.Cli) *cobra.Command {
 		Short: "Create a new builder instance",
 		Args:  cli.RequiresMaxArgs(1),
 		RunE: func(cmd *cobra.Command, args []string) error {
+			options.timeout = rootOpts.timeout
 			return runCreate(cmd.Context(), dockerCli, options, args)
 		},
 		ValidArgsFunction: completion.Disable,
diff --git a/commands/inspect.go b/commands/inspect.go
@@ -23,6 +23,7 @@ import (
 type inspectOptions struct {
 	bootstrap bool
 	builder   string
+	timeout   time.Duration
 }
 
 func runInspect(ctx context.Context, dockerCli command.Cli, in inspectOptions) error {
@@ -34,7 +35,7 @@ func runInspect(ctx context.Context, dockerCli command.Cli, in inspectOptions) e
 		return err
 	}
 
-	timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
+	timeoutCtx, cancel := context.WithTimeout(ctx, in.timeout)
 	defer cancel()
 
 	nodes, err := b.LoadNodes(timeoutCtx, builder.WithData())
@@ -147,6 +148,7 @@ func inspectCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 			if len(args) > 0 {
 				options.builder = args[0]
 			}
+			options.timeout = rootOpts.timeout
 			return runInspect(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
diff --git a/commands/ls.go b/commands/ls.go
@@ -35,7 +35,8 @@ const (
 )
 
 type lsOptions struct {
-	format string
+	format  string
+	timeout time.Duration
 }
 
 func runLs(ctx context.Context, dockerCli command.Cli, in lsOptions) error {
@@ -55,7 +56,7 @@ func runLs(ctx context.Context, dockerCli command.Cli, in lsOptions) error {
 		return err
 	}
 
-	timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
+	timeoutCtx, cancel := context.WithTimeout(ctx, in.timeout)
 	defer cancel()
 
 	eg, _ := errgroup.WithContext(timeoutCtx)
@@ -92,14 +93,15 @@ func runLs(ctx context.Context, dockerCli command.Cli, in lsOptions) error {
 	return nil
 }
 
-func lsCmd(dockerCli command.Cli) *cobra.Command {
+func lsCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 	var options lsOptions
 
 	cmd := &cobra.Command{
 		Use:   "ls",
 		Short: "List builder instances",
 		Args:  cli.ExactArgs(0),
 		RunE: func(cmd *cobra.Command, args []string) error {
+			options.timeout = rootOpts.timeout
 			return runLs(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.Disable,
diff --git a/commands/rm.go b/commands/rm.go
@@ -21,6 +21,7 @@ type rmOptions struct {
 	keepDaemon  bool
 	allInactive bool
 	force       bool
+	timeout     time.Duration
 }
 
 const (
@@ -109,6 +110,7 @@ func rmCmd(dockerCli command.Cli, rootOpts *rootOptions) *cobra.Command {
 				}
 				options.builders = args
 			}
+			options.timeout = rootOpts.timeout
 			return runRm(cmd.Context(), dockerCli, options)
 		},
 		ValidArgsFunction: completion.BuilderNames(dockerCli),
@@ -150,7 +152,7 @@ func rmAllInactive(ctx context.Context, txn *store.Txn, dockerCli command.Cli, i
 		return err
 	}
 
-	timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)
+	timeoutCtx, cancel := context.WithTimeout(ctx, in.timeout)
 	defer cancel()
 
 	eg, _ := errgroup.WithContext(timeoutCtx)
diff --git a/commands/root.go b/commands/root.go
@@ -2,6 +2,7 @@ package commands
 
 import (
 	"os"
+	"time"
 
 	debugcmd "github.com/docker/buildx/commands/debug"
 	imagetoolscmd "github.com/docker/buildx/commands/imagetools"
@@ -20,6 +21,8 @@ import (
 	"github.com/spf13/pflag"
 )
 
+const defaultTimeoutCli = 20 * time.Second
+
 func NewRootCmd(name string, isPlugin bool, dockerCli command.Cli) *cobra.Command {
 	cmd := &cobra.Command{
 		Short: "Docker Buildx",
@@ -74,6 +77,7 @@ func NewRootCmd(name string, isPlugin bool, dockerCli command.Cli) *cobra.Comman
 
 type rootOptions struct {
 	builder string
+	timeout time.Duration
 }
 
 func addCommands(cmd *cobra.Command, dockerCli command.Cli) {
@@ -83,10 +87,10 @@ func addCommands(cmd *cobra.Command, dockerCli command.Cli) {
 	cmd.AddCommand(
 		buildCmd(dockerCli, opts, nil),
 		bakeCmd(dockerCli, opts),
-		createCmd(dockerCli),
+		createCmd(dockerCli, opts),
 		dialStdioCmd(dockerCli, opts),
 		rmCmd(dockerCli, opts),
-		lsCmd(dockerCli),
+		lsCmd(dockerCli, opts),
 		useCmd(dockerCli, opts),
 		inspectCmd(dockerCli, opts),
 		stopCmd(dockerCli, opts),
@@ -112,4 +116,14 @@ func addCommands(cmd *cobra.Command, dockerCli command.Cli) {
 
 func rootFlags(options *rootOptions, flags *pflag.FlagSet) {
 	flags.StringVar(&options.builder, "builder", os.Getenv("BUILDX_BUILDER"), "Override the configured builder instance")
+
+	var timeoutDuration = defaultTimeoutCli
+	if value, ok := os.LookupEnv("BUILDX_TIMEOUT"); ok {
+		var err error
+		timeoutDuration, err = time.ParseDuration(value)
+		if err != nil {
+			timeoutDuration = defaultTimeoutCli
+		}
+	}
+	flags.DurationVar(&options.timeout, "timeout", timeoutDuration, "Override the default global timeout (20 seconds)")
 }
diff --git a/controller/control/controller.go b/controller/control/controller.go
@@ -3,6 +3,7 @@ package control
 import (
 	"context"
 	"io"
+	"time"
 
 	controllerapi "github.com/docker/buildx/controller/pb"
 	"github.com/docker/buildx/util/progress"
@@ -29,4 +30,5 @@ type ControlOptions struct {
 	ServerConfig string
 	Root         string
 	Detach       bool
+	Timeout      time.Duration
 }
diff --git a/controller/remote/controller.go b/controller/remote/controller.go
@@ -13,7 +13,6 @@ import (
 	"path/filepath"
 	"strconv"
 	"syscall"
-	"time"
 
 	"github.com/containerd/log"
 	"github.com/docker/buildx/build"
@@ -62,7 +61,7 @@ func NewRemoteBuildxController(ctx context.Context, dockerCli command.Cli, opts
 	serverRoot := filepath.Join(rootDir, "shared")
 
 	// connect to buildx server if it is already running
-	ctx2, cancel := context.WithTimeout(ctx, 1*time.Second)
+	ctx2, cancel := context.WithTimeout(ctx, opts.Timeout)
 	c, err := newBuildxClientAndCheck(ctx2, filepath.Join(serverRoot, defaultSocketFilename))
 	cancel()
 	if err != nil {
@@ -90,7 +89,7 @@ func NewRemoteBuildxController(ctx context.Context, dockerCli command.Cli, opts
 		go wait()
 
 		// wait for buildx server to be ready
-		ctx2, cancel = context.WithTimeout(ctx, 10*time.Second)
+		ctx2, cancel = context.WithTimeout(ctx, opts.Timeout)
 		c, err = newBuildxClientAndCheck(ctx2, filepath.Join(serverRoot, defaultSocketFilename))
 		cancel()
 		if err != nil {
diff --git a/driver/docker-container/driver.go b/driver/docker-container/driver.go
@@ -214,7 +214,7 @@ func (d *Driver) wait(ctx context.Context, l progress.SubLogger) error {
 			select {
 			case <-ctx.Done():
 				return ctx.Err()
-			case <-time.After(time.Duration(try*120) * time.Millisecond):
+			case <-time.After(d.Timeout / time.Microsecond):
 				try++
 				continue
 			}
diff --git a/driver/kubernetes/driver.go b/driver/kubernetes/driver.go
@@ -51,7 +51,6 @@ type Driver struct {
 	configMapClient  clientcorev1.ConfigMapInterface
 	podChooser       podchooser.PodChooser
 	defaultLoad      bool
-	timeout          time.Duration
 }
 
 func (d *Driver) IsMobyDriver() bool {
@@ -90,7 +89,7 @@ func (d *Driver) Bootstrap(ctx context.Context, l progress.Logger) error {
 			}
 		}
 		return sub.Wrap(
-			fmt.Sprintf("waiting for %d pods to be ready, timeout: %s", d.minReplicas, units.HumanDuration(d.timeout)),
+			fmt.Sprintf("waiting for %d pods to be ready, timeout: %s", d.minReplicas, units.HumanDuration(d.Timeout)),
 			func() error {
 				return d.wait(ctx)
 			})
@@ -104,8 +103,11 @@ func (d *Driver) wait(ctx context.Context) error {
 		depl *appsv1.Deployment
 	)
 
-	timeoutChan := time.After(d.timeout)
-	ticker := time.NewTicker(100 * time.Millisecond)
+	// 120 * 1000 * 1000 * 1000 =
+	// 100 * 1000 * 1000 =
+
+	timeoutChan := time.After(d.Timeout)
+	ticker := time.NewTicker(d.Timeout / time.Microsecond)
 	defer ticker.Stop()
 
 	for {
diff --git a/driver/kubernetes/factory.go b/driver/kubernetes/factory.go
@@ -20,7 +20,6 @@ import (
 const (
 	prioritySupported   = 40
 	priorityUnsupported = 80
-	defaultTimeout      = 120 * time.Second
 )
 
 func init() {
@@ -78,7 +77,7 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
 	}
 
 	d.defaultLoad = defaultLoad
-	d.timeout = timeout
+	d.Timeout = timeout
 
 	d.deployment, d.configMaps, err = manifest.NewDeployment(deploymentOpt)
 	if err != nil {
@@ -119,7 +118,7 @@ func (f *factory) processDriverOpts(deploymentName string, namespace string, cfg
 	}
 
 	defaultLoad := false
-	timeout := defaultTimeout
+	timeout := cfg.Timeout
 
 	deploymentOpt.Qemu.Image = bkimage.QemuImage
 
diff --git a/driver/manager.go b/driver/manager.go
@@ -6,6 +6,7 @@ import (
 	"sort"
 	"strings"
 	"sync"
+	"time"
 
 	dockerclient "github.com/docker/docker/client"
 	"github.com/moby/buildkit/client"
@@ -60,8 +61,11 @@ type InitConfig struct {
 	Platforms        []specs.Platform
 	ContextPathHash  string // can be used for determining pods in the driver instance
 	DialMeta         map[string][]string
+	Timeout          time.Duration
 }
 
+const defaultTimeoutDriver = 120 * time.Second
+
 var drivers map[string]Factory
 
 func Register(f Factory) {
@@ -117,6 +121,7 @@ func GetDriver(ctx context.Context, name string, f Factory, endpointAddr string,
 		ContextPathHash:  contextPathHash,
 		DialMeta:         dialMeta,
 		Files:            files,
+		Timeout:          defaultTimeoutDriver,
 	}
 	if f == nil {
 		var err error
diff --git a/driver/remote/driver.go b/driver/remote/driver.go
@@ -8,7 +8,6 @@ import (
 	"os"
 	"strings"
 	"sync"
-	"time"
 
 	"github.com/docker/buildx/driver"
 	util "github.com/docker/buildx/driver/remote/util"
@@ -47,7 +46,7 @@ func (d *Driver) Bootstrap(ctx context.Context, l progress.Logger) error {
 		return err
 	}
 	return progress.Wrap("[internal] waiting for connection", l, func(_ progress.SubLogger) error {
-		ctx, cancel := context.WithTimeout(ctx, 20*time.Second)
+		ctx, cancel := context.WithTimeout(ctx, d.Timeout)
 		defer cancel()
 		return c.Wait(ctx)
 	})
diff --git a/driver/remote/factory.go b/driver/remote/factory.go
@@ -6,6 +6,7 @@ import (
 	"path/filepath"
 	"strconv"
 	"strings"
+	"time"
 
 	// import connhelpers for special url schemes
 	_ "github.com/moby/buildkit/client/connhelper/dockercontainer"
@@ -87,6 +88,12 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
 				return nil, err
 			}
 			d.defaultLoad = parsed
+		case "timeout":
+			parsed, err := time.ParseDuration(v)
+			if err != nil {
+				return nil, err
+			}
+			d.Timeout = parsed
 		default:
 			return nil, errors.Errorf("invalid driver option %s for remote driver", k)
 		}

Original file line number	Diff line number	Diff line change
`@@ -343,6 +343,7 @@ type CreateOpts struct {`
`343`	`343`	`Use bool`
`344`	`344`	`Endpoint string`
`345`	`345`	`Append bool`
	`346`	`+ Timeout time.Duration`
`346`	`347`	`}`
`347`	`348`
`348`	`349`	`func Create(ctx context.Context, txn store.Txn, dockerCli command.Cli, opts CreateOpts) (Builder, error) {`
`@@ -522,7 +523,7 @@ func Create(ctx context.Context, txn *store.Txn, dockerCli command.Cli, opts Cre`
`522`	`523`	`return nil, err`
`523`	`524`	`}`
`524`	`525`
`525`		`- timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)`
	`526`	`+ timeoutCtx, cancel := context.WithTimeout(ctx, opts.Timeout)`
`526`	`527`	`defer cancel()`
`527`	`528`
`528`	`529`	`nodes, err := b.LoadNodes(timeoutCtx, WithData())`
Original file line number	Diff line number	Diff line change
`@@ -104,6 +104,8 @@ type buildOptions struct {`
`104`	`104`	`control.ControlOptions`
`105`	`105`
`106`	`106`	`invokeConfig *invokeConfig`
	`107`	`+`
	`108`	`+ timeout time.Duration`
`107`	`109`	`}`
`108`	`110`
`109`	`111`	`func (o buildOptions) toControllerOptions() (controllerapi.BuildOptions, error) {`
Original file line number	Diff line number	Diff line change
`@@ -21,6 +21,7 @@ type rmOptions struct {`
`21`	`21`	`keepDaemon bool`
`22`	`22`	`allInactive bool`
`23`	`23`	`force bool`
	`24`	`+ timeout time.Duration`
`24`	`25`	`}`
`25`	`26`
`26`	`27`	`const (`
`@@ -109,6 +110,7 @@ func rmCmd(dockerCli command.Cli, rootOpts rootOptions) cobra.Command {`
`109`	`110`	`}`
`110`	`111`	`options.builders = args`
`111`	`112`	`}`
	`113`	`+ options.timeout = rootOpts.timeout`
`112`	`114`	`return runRm(cmd.Context(), dockerCli, options)`
`113`	`115`	`},`
`114`	`116`	`ValidArgsFunction: completion.BuilderNames(dockerCli),`
`@@ -150,7 +152,7 @@ func rmAllInactive(ctx context.Context, txn *store.Txn, dockerCli command.Cli, i`
`150`	`152`	`return err`
`151`	`153`	`}`
`152`	`154`
`153`		`- timeoutCtx, cancel := context.WithTimeout(ctx, 20*time.Second)`
	`155`	`+ timeoutCtx, cancel := context.WithTimeout(ctx, in.timeout)`
`154`	`156`	`defer cancel()`
`155`	`157`
`156`	`158`	`eg, _ := errgroup.WithContext(timeoutCtx)`
Original file line number	Diff line number	Diff line change
`@@ -214,7 +214,7 @@ func (d *Driver) wait(ctx context.Context, l progress.SubLogger) error {`
`214`	`214`	`select {`
`215`	`215`	`case <-ctx.Done():`
`216`	`216`	`return ctx.Err()`
`217`		`- case <-time.After(time.Duration(try120) time.Millisecond):`
	`217`	`+ case <-time.After(d.Timeout / time.Microsecond):`
`218`	`218`	`try++`
`219`	`219`	`continue`
`220`	`220`	`}`