From 9822409b672d5ae6c655112526e865288e830be0 Mon Sep 17 00:00:00 2001
From: CrazyMax <1951866+crazy-max@users.noreply.github.com>
Date: Tue, 6 Feb 2024 14:33:36 +0100
Subject: [PATCH] docker-container: restart-policy opt

Signed-off-by: CrazyMax <1951866+crazy-max@users.noreply.github.com>
---
 driver/docker-container/driver.go  | 26 ++++++++++++++------------
 driver/docker-container/factory.go | 17 ++++++++++++++++-
 tests/create.go                    | 20 ++++++++++++++++++++
 3 files changed, 50 insertions(+), 13 deletions(-)

diff --git a/driver/docker-container/driver.go b/driver/docker-container/driver.go
index 4998b0cc..2f36ff75 100644
--- a/driver/docker-container/driver.go
+++ b/driver/docker-container/driver.go
@@ -44,17 +44,18 @@ type Driver struct {
 
 	// if you add fields, remember to update docs:
 	// https://github.com/docker/docs/blob/main/content/build/drivers/docker-container.md
-	netMode      string
-	image        string
-	memory       opts.MemBytes
-	memorySwap   opts.MemSwapBytes
-	cpuQuota     int64
-	cpuPeriod    int64
-	cpuShares    int64
-	cpusetCpus   string
-	cpusetMems   string
-	cgroupParent string
-	env          []string
+	netMode       string
+	image         string
+	memory        opts.MemBytes
+	memorySwap    opts.MemSwapBytes
+	cpuQuota      int64
+	cpuPeriod     int64
+	cpuShares     int64
+	cpusetCpus    string
+	cpusetMems    string
+	cgroupParent  string
+	restartPolicy container.RestartPolicy
+	env           []string
 }
 
 func (d *Driver) IsMobyDriver() bool {
@@ -121,7 +122,8 @@ func (d *Driver) create(ctx context.Context, l progress.SubLogger) error {
 	useInit := true // let it cleanup exited processes created by BuildKit's container API
 	return l.Wrap("creating container "+d.Name, func() error {
 		hc := &container.HostConfig{
-			Privileged: true,
+			Privileged:    true,
+			RestartPolicy: d.restartPolicy,
 			Mounts: []mount.Mount{
 				{
 					Type:   mount.TypeVolume,
diff --git a/driver/docker-container/factory.go b/driver/docker-container/factory.go
index 00564be7..ff69ae0b 100644
--- a/driver/docker-container/factory.go
+++ b/driver/docker-container/factory.go
@@ -7,12 +7,14 @@ import (
 	"strings"
 
 	"github.com/docker/buildx/driver"
+	dockeropts "github.com/docker/cli/opts"
 	dockerclient "github.com/docker/docker/client"
 	"github.com/pkg/errors"
 )
 
 const prioritySupported = 30
 const priorityUnsupported = 70
+const defaultRestartPolicy = "unless-stopped"
 
 func init() {
 	driver.Register(&factory{})
@@ -40,7 +42,15 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
 	if cfg.DockerAPI == nil {
 		return nil, errors.Errorf("%s driver requires docker API access", f.Name())
 	}
-	d := &Driver{factory: f, InitConfig: cfg}
+	rp, err := dockeropts.ParseRestartPolicy(defaultRestartPolicy)
+	if err != nil {
+		return nil, err
+	}
+	d := &Driver{
+		factory:       f,
+		InitConfig:    cfg,
+		restartPolicy: rp,
+	}
 	for k, v := range cfg.DriverOpts {
 		switch {
 		case k == "network":
@@ -82,6 +92,11 @@ func (f *factory) New(ctx context.Context, cfg driver.InitConfig) (driver.Driver
 			d.cpusetMems = v
 		case k == "cgroup-parent":
 			d.cgroupParent = v
+		case k == "restart-policy":
+			d.restartPolicy, err = dockeropts.ParseRestartPolicy(v)
+			if err != nil {
+				return nil, err
+			}
 		case strings.HasPrefix(k, "env."):
 			envName := strings.TrimPrefix(k, "env.")
 			if envName == "" {
diff --git a/tests/create.go b/tests/create.go
index 65476c24..a37fa3da 100644
--- a/tests/create.go
+++ b/tests/create.go
@@ -17,6 +17,7 @@ func createCmd(sb integration.Sandbox, opts ...cmdOpt) (string, error) {
 
 var createTests = []func(t *testing.T, sb integration.Sandbox){
 	testCreateMemoryLimit,
+	testCreateRestartAlways,
 }
 
 func testCreateMemoryLimit(t *testing.T, sb integration.Sandbox) {
@@ -37,3 +38,22 @@ func testCreateMemoryLimit(t *testing.T, sb integration.Sandbox) {
 	require.NoError(t, err, out)
 	builderName = strings.TrimSpace(out)
 }
+
+func testCreateRestartAlways(t *testing.T, sb integration.Sandbox) {
+	if sb.Name() != "docker-container" {
+		t.Skip("only testing for docker-container driver")
+	}
+
+	var builderName string
+	t.Cleanup(func() {
+		if builderName == "" {
+			return
+		}
+		out, err := rmCmd(sb, withArgs(builderName))
+		require.NoError(t, err, out)
+	})
+
+	out, err := createCmd(sb, withArgs("--driver", "docker-container", "--driver-opt", "restart-policy=always"))
+	require.NoError(t, err, out)
+	builderName = strings.TrimSpace(out)
+}