From edd76bfd707a790c2e01aefceebf225eaa0b5b47 Mon Sep 17 00:00:00 2001
From: Ulysses Souza <ulyssessouza@gmail.com>
Date: Mon, 26 Jun 2023 15:52:50 +0200
Subject: [PATCH] Add `docker compose wait`

Signed-off-by: Ulysses Souza <ulyssessouza@gmail.com>
---
 cmd/compose/compose.go                  |  1 +
 cmd/compose/wait.go                     | 72 +++++++++++++++++++++++++
 docs/reference/compose.md               |  1 +
 docs/reference/compose_wait.md          | 15 ++++++
 docs/reference/docker_compose.yaml      |  2 +
 docs/reference/docker_compose_wait.yaml | 34 ++++++++++++
 pkg/api/api.go                          |  9 ++++
 pkg/api/proxy.go                        | 12 ++++-
 pkg/compose/wait.go                     | 67 +++++++++++++++++++++++
 pkg/e2e/fixtures/wait/compose.yaml      | 11 ++++
 pkg/e2e/wait_test.go                    | 72 +++++++++++++++++++++++++
 pkg/mocks/mock_docker_compose_api.go    | 15 ++++++
 12 files changed, 310 insertions(+), 1 deletion(-)
 create mode 100644 cmd/compose/wait.go
 create mode 100644 docs/reference/compose_wait.md
 create mode 100644 docs/reference/docker_compose_wait.yaml
 create mode 100644 pkg/compose/wait.go
 create mode 100644 pkg/e2e/fixtures/wait/compose.yaml
 create mode 100644 pkg/e2e/wait_test.go

diff --git a/cmd/compose/compose.go b/cmd/compose/compose.go
index 0a2fc1fa8..04f4ea88d 100644
--- a/cmd/compose/compose.go
+++ b/cmd/compose/compose.go
@@ -431,6 +431,7 @@ func RootCommand(streams command.Cli, backend api.Service) *cobra.Command { //no
 		pullCommand(&opts, backend),
 		createCommand(&opts, backend),
 		copyCommand(&opts, backend),
+		waitCommand(&opts, backend),
 		alphaCommand(&opts, backend),
 	)
 
diff --git a/cmd/compose/wait.go b/cmd/compose/wait.go
new file mode 100644
index 000000000..6570b2d25
--- /dev/null
+++ b/cmd/compose/wait.go
@@ -0,0 +1,72 @@
+/*
+   Copyright 2023 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package compose
+
+import (
+	"context"
+	"os"
+
+	"github.com/docker/cli/cli"
+	"github.com/docker/compose/v2/pkg/api"
+	"github.com/spf13/cobra"
+)
+
+type waitOptions struct {
+	*ProjectOptions
+
+	services []string
+
+	downProject bool
+}
+
+func waitCommand(p *ProjectOptions, backend api.Service) *cobra.Command {
+	opts := waitOptions{
+		ProjectOptions: p,
+	}
+
+	var statusCode int64
+	var err error
+	cmd := &cobra.Command{
+		Use:   "wait SERVICE [SERVICE...] [OPTIONS]",
+		Short: "Block until the first service container stops",
+		Args:  cli.RequiresMinArgs(1),
+		RunE: Adapt(func(ctx context.Context, services []string) error {
+			opts.services = services
+			statusCode, err = runWait(ctx, backend, &opts)
+			return err
+		}),
+		PostRun: func(cmd *cobra.Command, args []string) {
+			os.Exit(int(statusCode))
+		},
+	}
+
+	cmd.Flags().BoolVar(&opts.downProject, "down-project", false, "Drops project when the first container stops")
+
+	return cmd
+}
+
+func runWait(ctx context.Context, backend api.Service, opts *waitOptions) (int64, error) {
+	_, name, err := opts.projectOrName()
+	if err != nil {
+		return 0, err
+	}
+
+	return backend.Wait(ctx, name, api.WaitOptions{
+		Services:                   opts.services,
+		DownProjectOnContainerExit: opts.downProject,
+	})
+}
diff --git a/docs/reference/compose.md b/docs/reference/compose.md
index 337ac3f14..0c7e14c39 100644
--- a/docs/reference/compose.md
+++ b/docs/reference/compose.md
@@ -33,6 +33,7 @@ Define and run multi-container applications with Docker.
 | [`unpause`](compose_unpause.md) | Unpause services                                                        |
 | [`up`](compose_up.md)           | Create and start containers                                             |
 | [`version`](compose_version.md) | Show the Docker Compose version information                             |
+| [`wait`](compose_wait.md)       | Block until the first service container stops                           |
 
 
 ### Options
diff --git a/docs/reference/compose_wait.md b/docs/reference/compose_wait.md
new file mode 100644
index 000000000..9c9ff6f1c
--- /dev/null
+++ b/docs/reference/compose_wait.md
@@ -0,0 +1,15 @@
+# docker compose wait
+
+<!---MARKER_GEN_START-->
+Block until the first service container stops
+
+### Options
+
+| Name             | Type | Default | Description                                  |
+|:-----------------|:-----|:--------|:---------------------------------------------|
+| `--down-project` |      |         | Drops project when the first container stops |
+| `--dry-run`      |      |         | Execute command in dry run mode              |
+
+
+<!---MARKER_GEN_END-->
+
diff --git a/docs/reference/docker_compose.yaml b/docs/reference/docker_compose.yaml
index 1f3a49ebd..984d56888 100644
--- a/docs/reference/docker_compose.yaml
+++ b/docs/reference/docker_compose.yaml
@@ -171,6 +171,7 @@ cname:
     - docker compose unpause
     - docker compose up
     - docker compose version
+    - docker compose wait
 clink:
     - docker_compose_build.yaml
     - docker_compose_config.yaml
@@ -197,6 +198,7 @@ clink:
     - docker_compose_unpause.yaml
     - docker_compose_up.yaml
     - docker_compose_version.yaml
+    - docker_compose_wait.yaml
 options:
     - option: ansi
       value_type: string
diff --git a/docs/reference/docker_compose_wait.yaml b/docs/reference/docker_compose_wait.yaml
new file mode 100644
index 000000000..8e5ff7eb4
--- /dev/null
+++ b/docs/reference/docker_compose_wait.yaml
@@ -0,0 +1,34 @@
+command: docker compose wait
+short: Block until the first service container stops
+long: Block until the first service container stops
+usage: docker compose wait SERVICE [SERVICE...] [OPTIONS]
+pname: docker compose
+plink: docker_compose.yaml
+options:
+    - option: down-project
+      value_type: bool
+      default_value: "false"
+      description: Drops project when the first container stops
+      deprecated: false
+      hidden: false
+      experimental: false
+      experimentalcli: false
+      kubernetes: false
+      swarm: false
+inherited_options:
+    - option: dry-run
+      value_type: bool
+      default_value: "false"
+      description: Execute command in dry run mode
+      deprecated: false
+      hidden: false
+      experimental: false
+      experimentalcli: false
+      kubernetes: false
+      swarm: false
+deprecated: false
+experimental: false
+experimentalcli: false
+kubernetes: false
+swarm: false
+
diff --git a/pkg/api/api.go b/pkg/api/api.go
index e1954be39..e6ae1a961 100644
--- a/pkg/api/api.go
+++ b/pkg/api/api.go
@@ -84,6 +84,15 @@ type Service interface {
 	Watch(ctx context.Context, project *types.Project, services []string, options WatchOptions) error
 	// Viz generates a graphviz graph of the project services
 	Viz(ctx context.Context, project *types.Project, options VizOptions) (string, error)
+	// Wait blocks until at least one of the services' container exits
+	Wait(ctx context.Context, projectName string, options WaitOptions) (int64, error)
+}
+
+type WaitOptions struct {
+	// Services passed in the command line to be waited
+	Services []string
+	// Executes a down when a container exits
+	DownProjectOnContainerExit bool
 }
 
 type VizOptions struct {
diff --git a/pkg/api/proxy.go b/pkg/api/proxy.go
index 30f9aa273..aeb773c30 100644
--- a/pkg/api/proxy.go
+++ b/pkg/api/proxy.go
@@ -54,6 +54,7 @@ type ServiceProxy struct {
 	MaxConcurrencyFn     func(parallel int)
 	DryRunModeFn         func(ctx context.Context, dryRun bool) (context.Context, error)
 	VizFn                func(ctx context.Context, project *types.Project, options VizOptions) (string, error)
+	WaitFn               func(ctx context.Context, projectName string, options WaitOptions) (int64, error)
 	interceptors         []Interceptor
 }
 
@@ -95,6 +96,7 @@ func (s *ServiceProxy) WithService(service Service) *ServiceProxy {
 	s.MaxConcurrencyFn = service.MaxConcurrency
 	s.DryRunModeFn = service.DryRunMode
 	s.VizFn = service.Viz
+	s.WaitFn = service.Wait
 	return s
 }
 
@@ -325,7 +327,7 @@ func (s *ServiceProxy) Watch(ctx context.Context, project *types.Project, servic
 	return s.WatchFn(ctx, project, services, options)
 }
 
-// Viz implements Viz interface
+// Viz implements Service interface
 func (s *ServiceProxy) Viz(ctx context.Context, project *types.Project, options VizOptions) (string, error) {
 	if s.VizFn == nil {
 		return "", ErrNotImplemented
@@ -333,6 +335,14 @@ func (s *ServiceProxy) Viz(ctx context.Context, project *types.Project, options
 	return s.VizFn(ctx, project, options)
 }
 
+// Wait implements Service interface
+func (s *ServiceProxy) Wait(ctx context.Context, projectName string, options WaitOptions) (int64, error) {
+	if s.WaitFn == nil {
+		return 0, ErrNotImplemented
+	}
+	return s.WaitFn(ctx, projectName, options)
+}
+
 func (s *ServiceProxy) MaxConcurrency(i int) {
 	s.MaxConcurrencyFn(i)
 }
diff --git a/pkg/compose/wait.go b/pkg/compose/wait.go
new file mode 100644
index 000000000..952e65cf8
--- /dev/null
+++ b/pkg/compose/wait.go
@@ -0,0 +1,67 @@
+/*
+   Copyright 2020 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package compose
+
+import (
+	"context"
+	"fmt"
+
+	"github.com/docker/compose/v2/pkg/api"
+	"golang.org/x/sync/errgroup"
+)
+
+func (s *composeService) Wait(ctx context.Context, projectName string, options api.WaitOptions) (int64, error) {
+	containers, err := s.getContainers(ctx, projectName, oneOffInclude, false, options.Services...)
+	if err != nil {
+		return 0, err
+	}
+	if len(containers) == 0 {
+		return 0, fmt.Errorf("no containers for project %q", projectName)
+	}
+
+	eg, waitCtx := errgroup.WithContext(ctx)
+	var statusCode int64
+	for _, c := range containers {
+		c := c
+		eg.Go(func() error {
+			var err error
+			resultC, errC := s.dockerCli.Client().ContainerWait(waitCtx, c.ID, "")
+
+			select {
+			case result := <-resultC:
+				fmt.Fprintf(s.dockerCli.Out(), "container %q exited with status code %d\n", c.ID, result.StatusCode)
+				statusCode = result.StatusCode
+			case err = <-errC:
+			}
+
+			return err
+		})
+	}
+
+	err = eg.Wait()
+	if err != nil {
+		return 42, err // Ignore abort flag in case of error in wait
+	}
+
+	if options.DownProjectOnContainerExit {
+		return statusCode, s.Down(ctx, projectName, api.DownOptions{
+			RemoveOrphans: true,
+		})
+	}
+
+	return statusCode, err
+}
diff --git a/pkg/e2e/fixtures/wait/compose.yaml b/pkg/e2e/fixtures/wait/compose.yaml
new file mode 100644
index 000000000..1a001e6fa
--- /dev/null
+++ b/pkg/e2e/fixtures/wait/compose.yaml
@@ -0,0 +1,11 @@
+services:
+  faster:
+    image: alpine
+    command: sleep 2
+  slower:
+    image: alpine
+    command: sleep 5
+  infinity:
+    image: alpine
+    command: sleep infinity
+
diff --git a/pkg/e2e/wait_test.go b/pkg/e2e/wait_test.go
new file mode 100644
index 000000000..f607f5ea5
--- /dev/null
+++ b/pkg/e2e/wait_test.go
@@ -0,0 +1,72 @@
+/*
+   Copyright 2020 Docker Compose CLI authors
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.
+*/
+
+package e2e
+
+import (
+	"strings"
+	"testing"
+	"time"
+
+	"gotest.tools/v3/assert"
+)
+
+func TestWaitOnFaster(t *testing.T) {
+	const projectName = "e2e-wait-faster"
+	c := NewParallelCLI(t)
+
+	c.RunDockerComposeCmd(t, "-f", "./fixtures/wait/compose.yaml", "--project-name", projectName, "up", "-d")
+	c.RunDockerComposeCmd(t, "--project-name", projectName, "wait", "faster")
+}
+
+func TestWaitOnSlower(t *testing.T) {
+	const projectName = "e2e-wait-slower"
+	c := NewParallelCLI(t)
+
+	c.RunDockerComposeCmd(t, "-f", "./fixtures/wait/compose.yaml", "--project-name", projectName, "up", "-d")
+	c.RunDockerComposeCmd(t, "--project-name", projectName, "wait", "slower")
+}
+
+func TestWaitOnInfinity(t *testing.T) {
+	const projectName = "e2e-wait-infinity"
+	c := NewParallelCLI(t)
+
+	c.RunDockerComposeCmd(t, "-f", "./fixtures/wait/compose.yaml", "--project-name", projectName, "up", "-d")
+
+	finished := make(chan struct{})
+	ticker := time.NewTicker(7 * time.Second)
+	go func() {
+		c.RunDockerComposeCmd(t, "--project-name", projectName, "wait", "infinity")
+		finished <- struct{}{}
+	}()
+
+	select {
+	case <-finished:
+		t.Fatal("wait infinity should not finish")
+	case <-ticker.C:
+	}
+}
+
+func TestWaitAndDrop(t *testing.T) {
+	const projectName = "e2e-wait-and-drop"
+	c := NewParallelCLI(t)
+
+	c.RunDockerComposeCmd(t, "-f", "./fixtures/wait/compose.yaml", "--project-name", projectName, "up", "-d")
+	c.RunDockerComposeCmd(t, "--project-name", projectName, "wait", "--down-project", "faster")
+
+	res := c.RunDockerCmd(t, "ps", "--all")
+	assert.Assert(t, !strings.Contains(res.Combined(), projectName), res.Combined())
+}
diff --git a/pkg/mocks/mock_docker_compose_api.go b/pkg/mocks/mock_docker_compose_api.go
index 120affd98..ffb62f3fb 100644
--- a/pkg/mocks/mock_docker_compose_api.go
+++ b/pkg/mocks/mock_docker_compose_api.go
@@ -423,6 +423,21 @@ func (mr *MockServiceMockRecorder) Viz(ctx, project, options interface{}) *gomoc
 	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Viz", reflect.TypeOf((*MockService)(nil).Viz), ctx, project, options)
 }
 
+// Wait mocks base method.
+func (m *MockService) Wait(ctx context.Context, projectName string, options api.WaitOptions) (int64, error) {
+	m.ctrl.T.Helper()
+	ret := m.ctrl.Call(m, "Wait", ctx, projectName, options)
+	ret0, _ := ret[0].(int64)
+	ret1, _ := ret[1].(error)
+	return ret0, ret1
+}
+
+// Wait indicates an expected call of Wait.
+func (mr *MockServiceMockRecorder) Wait(ctx, projectName, options interface{}) *gomock.Call {
+	mr.mock.ctrl.T.Helper()
+	return mr.mock.ctrl.RecordCallWithMethodType(mr.mock, "Wait", reflect.TypeOf((*MockService)(nil).Wait), ctx, projectName, options)
+}
+
 // Watch mocks base method.
 func (m *MockService) Watch(ctx context.Context, project *types.Project, services []string, options api.WatchOptions) error {
 	m.ctrl.T.Helper()