Pull busybox image if we need to

skip-mev · Jan 31, 2024 · 9ecefbe · 9ecefbe
1 parent 9eb28e4
commit 9ecefbe
Show file tree

Hide file tree

Showing 3 changed files with 38 additions and 24 deletions.
diff --git a/provider/docker/docker_provider.go b/provider/docker/docker_provider.go
@@ -26,6 +26,7 @@ type Provider struct {
 	dockerNetworkName string
 	networkMu         sync.RWMutex
 	listeners         map[string]Listeners
+	builderImageName  string
 }
 
 func NewDockerProvider(ctx context.Context, logger *zap.Logger, providerName string, dockerOpts ...client.Opt) (*Provider, error) {
@@ -41,10 +42,11 @@ func NewDockerProvider(ctx context.Context, logger *zap.Logger, providerName str
 	}
 
 	dockerProvider := &Provider{
-		dockerClient: dockerClient,
-		listeners:    map[string]Listeners{},
-		logger:       logger.Named("docker_provider"),
-		name:         providerName,
+		dockerClient:     dockerClient,
+		listeners:        map[string]Listeners{},
+		logger:           logger.Named("docker_provider"),
+		name:             providerName,
+		builderImageName: "busybox:latest",
 	}
 
 	dockerProvider.dockerNetworkName = fmt.Sprintf("petri-network-%s", util.RandomString(5))

diff --git a/provider/docker/task.go b/provider/docker/task.go
@@ -23,12 +23,8 @@ import (
 func (p *Provider) CreateTask(ctx context.Context, logger *zap.Logger, definition provider.TaskDefinition) (string, error) {
 	logger = logger.Named("docker_provider")
 
-	_, _, err := p.dockerClient.ImageInspectWithRaw(ctx, definition.Image.Image)
-	if err != nil {
-		logger.Info("image not found, pulling", zap.String("image", definition.Image.Image))
-		if err := p.pullImage(ctx, definition); err != nil {
-			return "", err
-		}
+	if err := p.pullImage(ctx, definition.Image.Image); err != nil {
+		return "", err
 	}
 
 	portSet := convertTaskDefinitionPortsToPortSet(definition)
@@ -99,18 +95,18 @@ func (p *Provider) CreateTask(ctx context.Context, logger *zap.Logger, definitio
 	return createdContainer.ID, nil
 }
 
-func (p *Provider) pullImage(ctx context.Context, definition provider.TaskDefinition) error {
-	p.logger.Info("pulling image", zap.String("image", definition.Image.Image))
-	resp, err := p.dockerClient.ImagePull(ctx, definition.Image.Image, types.ImagePullOptions{})
+func (p *Provider) pullImage(ctx context.Context, imageName string) error {
+	_, _, err := p.dockerClient.ImageInspectWithRaw(ctx, imageName)
 	if err != nil {
-		return err
-	}
-
-	defer resp.Close()
-	// throw away the image pull stdout response
-	_, err = io.Copy(io.Discard, resp)
+		p.logger.Info("image not found, pulling", zap.String("image", imageName))
+		resp, err := p.dockerClient.ImagePull(ctx, imageName, types.ImagePullOptions{})
+		if err != nil {
+			return err
+		}
+		defer resp.Close()
 
-	if err != nil {
+		// throw away the image pull stdout response
+		_, err = io.Copy(io.Discard, resp)
 		return err
 	}
 	return nil

diff --git a/provider/docker/volume.go b/provider/docker/volume.go
@@ -71,12 +71,16 @@ func (p *Provider) WriteFile(ctx context.Context, id, relPath string, content []
 
 	containerName := fmt.Sprintf("petri-writefile-%d", time.Now().UnixNano())
 
+	if err := p.pullImage(ctx, p.builderImageName); err != nil {
+		return err
+	}
+
 	logger.Debug("creating writefile container")
 
 	cc, err := p.dockerClient.ContainerCreate(
 		ctx,
 		&container.Config{
-			Image: "busybox:latest",
+			Image: p.builderImageName,
 
 			Entrypoint: []string{"sh", "-c"},
 			Cmd: []string{
@@ -202,12 +206,16 @@ func (p *Provider) ReadFile(ctx context.Context, id, relPath string) ([]byte, er
 
 	containerName := fmt.Sprintf("petri-getfile-%d", time.Now().UnixNano())
 
+	if err := p.pullImage(ctx, p.builderImageName); err != nil {
+		return nil, err
+	}
+
 	logger.Debug("creating getfile container")
 
 	cc, err := p.dockerClient.ContainerCreate(
 		ctx,
 		&container.Config{
-			Image: "busybox:latest",
+			Image: p.builderImageName,
 
 			Labels: map[string]string{
 				providerLabelName: p.name,
@@ -290,10 +298,14 @@ func (p *Provider) DownloadDir(ctx context.Context, id, relPath, localPath strin
 
 	logger.Debug("creating getdir container")
 
+	if err := p.pullImage(ctx, p.builderImageName); err != nil {
+		return err
+	}
+
 	cc, err := p.dockerClient.ContainerCreate(
 		ctx,
 		&container.Config{
-			Image: "busybox:latest",
+			Image: p.builderImageName,
 
 			Labels: map[string]string{
 				providerLabelName: p.name,
@@ -370,12 +382,16 @@ func (p *Provider) SetVolumeOwner(ctx context.Context, volumeName, uid, gid stri
 
 	containerName := fmt.Sprintf("petri-setowner-%d", time.Now().UnixNano())
 
+	if err := p.pullImage(ctx, p.builderImageName); err != nil {
+		return err
+	}
+
 	logger.Debug("creating volume-owner container")
 
 	cc, err := p.dockerClient.ContainerCreate(
 		ctx,
 		&container.Config{
-			Image:      "busybox:latest",
+			Image:      p.builderImageName,
 			Entrypoint: []string{"sh", "-c"},
 			Cmd: []string{
 				`chown "$2:$3" "$1" && chmod 0700 "$1"`,