From 7205d918adfb3e8b368094d040c875722487f7f9 Mon Sep 17 00:00:00 2001
From: Nicolas De Loof <nicolas.deloof@gmail.com>
Date: Mon, 15 Nov 2021 08:48:31 +0100
Subject: [PATCH] interrupt printer when `compose log` is cancelled

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>
---
 pkg/compose/logs.go    |  2 +-
 pkg/compose/printer.go | 84 ++++++++++++++++++++++--------------------
 pkg/compose/up.go      |  2 +-
 3 files changed, 47 insertions(+), 41 deletions(-)

diff --git a/pkg/compose/logs.go b/pkg/compose/logs.go
index e586130c..1b143a6b 100644
--- a/pkg/compose/logs.go
+++ b/pkg/compose/logs.go
@@ -42,7 +42,7 @@ func (s *composeService) Logs(ctx context.Context, projectName string, consumer
 			})
 		})
 		eg.Go(func() error {
-			_, err := printer.Run(false, "", nil)
+			_, err := printer.Run(ctx, false, "", nil)
 			return err
 		})
 	}
diff --git a/pkg/compose/printer.go b/pkg/compose/printer.go
index 122ef731..4816f70e 100644
--- a/pkg/compose/printer.go
+++ b/pkg/compose/printer.go
@@ -17,6 +17,7 @@
 package compose
 
 import (
+	"context"
 	"fmt"
 
 	"github.com/docker/compose/v2/pkg/api"
@@ -27,7 +28,7 @@ import (
 // logPrinter watch application containers an collect their logs
 type logPrinter interface {
 	HandleEvent(event api.ContainerEvent)
-	Run(cascadeStop bool, exitCodeFrom string, stopFn func() error) (int, error)
+	Run(ctx context.Context, cascadeStop bool, exitCodeFrom string, stopFn func() error) (int, error)
 	Cancel()
 }
 
@@ -56,56 +57,61 @@ func (p *printer) HandleEvent(event api.ContainerEvent) {
 	p.queue <- event
 }
 
-func (p *printer) Run(cascadeStop bool, exitCodeFrom string, stopFn func() error) (int, error) {
+//nolint:gocyclo
+func (p *printer) Run(ctx context.Context, cascadeStop bool, exitCodeFrom string, stopFn func() error) (int, error) {
 	var (
 		aborting bool
 		exitCode int
 	)
 	containers := map[string]struct{}{}
 	for {
-		event := <-p.queue
-		container := event.Container
-		switch event.Type {
-		case api.UserCancel:
-			aborting = true
-		case api.ContainerEventAttach:
-			if _, ok := containers[container]; ok {
-				continue
-			}
-			containers[container] = struct{}{}
-			p.consumer.Register(container)
-		case api.ContainerEventExit:
-			if !event.Restarting {
-				delete(containers, container)
-			}
-			if !aborting {
-				p.consumer.Status(container, fmt.Sprintf("exited with code %d", event.ExitCode))
-			}
-			if cascadeStop {
+		select {
+		case <-ctx.Done():
+			return exitCode, ctx.Err()
+		case event := <-p.queue:
+			container := event.Container
+			switch event.Type {
+			case api.UserCancel:
+				aborting = true
+			case api.ContainerEventAttach:
+				if _, ok := containers[container]; ok {
+					continue
+				}
+				containers[container] = struct{}{}
+				p.consumer.Register(container)
+			case api.ContainerEventExit:
+				if !event.Restarting {
+					delete(containers, container)
+				}
 				if !aborting {
-					aborting = true
-					fmt.Println("Aborting on container exit...")
-					err := stopFn()
-					if err != nil {
-						return 0, err
+					p.consumer.Status(container, fmt.Sprintf("exited with code %d", event.ExitCode))
+				}
+				if cascadeStop {
+					if !aborting {
+						aborting = true
+						fmt.Println("Aborting on container exit...")
+						err := stopFn()
+						if err != nil {
+							return 0, err
+						}
+					}
+					if exitCodeFrom == "" {
+						exitCodeFrom = event.Service
+					}
+					if exitCodeFrom == event.Service {
+						logrus.Error(event.ExitCode)
+						exitCode = event.ExitCode
 					}
 				}
-				if exitCodeFrom == "" {
-					exitCodeFrom = event.Service
+				if len(containers) == 0 {
+					// Last container terminated, done
+					return exitCode, nil
 				}
-				if exitCodeFrom == event.Service {
-					logrus.Error(event.ExitCode)
-					exitCode = event.ExitCode
+			case api.ContainerEventLog:
+				if !aborting {
+					p.consumer.Log(container, event.Service, event.Line)
 				}
 			}
-			if len(containers) == 0 {
-				// Last container terminated, done
-				return exitCode, nil
-			}
-		case api.ContainerEventLog:
-			if !aborting {
-				p.consumer.Log(container, event.Service, event.Line)
-			}
 		}
 	}
 }
diff --git a/pkg/compose/up.go b/pkg/compose/up.go
index 47256b7b..88d0df1f 100644
--- a/pkg/compose/up.go
+++ b/pkg/compose/up.go
@@ -80,7 +80,7 @@ func (s *composeService) Up(ctx context.Context, project *types.Project, options
 	var exitCode int
 	eg, ctx := errgroup.WithContext(ctx)
 	eg.Go(func() error {
-		code, err := printer.Run(options.Start.CascadeStop, options.Start.ExitCodeFrom, stopFunc)
+		code, err := printer.Run(context.Background(), options.Start.CascadeStop, options.Start.ExitCodeFrom, stopFunc)
 		exitCode = code
 		return err
 	})