From b7dd74058c8e6b4d7b318548ce36aec1cc8d7ddc Mon Sep 17 00:00:00 2001
From: Zero <tobewhatwewant@outlook.com>
Date: Tue, 23 Jan 2024 01:30:11 +0800
Subject: [PATCH] feat: support driver-opt restart-policy, default
 unless-stopped

---
 driver/docker-container/driver.go  | 26 ++++++++++++++------------
 driver/docker-container/factory.go | 19 ++++++++++++++++++-
 2 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/driver/docker-container/driver.go b/driver/docker-container/driver.go
index 4998b0cc008..2f36ff754f3 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 00564be7ab4..58030221e44 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,16 @@ 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 +93,12 @@ 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":
+			vv, err := dockeropts.ParseRestartPolicy(v)
+			if err != nil {
+				return nil, err
+			}
+			d.restartPolicy = vv
 		case strings.HasPrefix(k, "env."):
 			envName := strings.TrimPrefix(k, "env.")
 			if envName == "" {