From 1b43e0c1d64b94b60f6b6aafd0f443ea42eaaca7 Mon Sep 17 00:00:00 2001
From: Nicolas De Loof <nicolas.deloof@gmail.com>
Date: Fri, 21 Aug 2020 15:56:31 +0200
Subject: [PATCH] Poll ECS logs until context is canceled

Signed-off-by: Nicolas De Loof <nicolas.deloof@gmail.com>
---
 ecs/logs.go | 11 +----------
 ecs/sdk.go  | 45 +++++++++++++++++++++++++--------------------
 2 files changed, 26 insertions(+), 30 deletions(-)

diff --git a/ecs/logs.go b/ecs/logs.go
index 84e913ea6..716736639 100644
--- a/ecs/logs.go
+++ b/ecs/logs.go
@@ -21,8 +21,6 @@ import (
 	"context"
 	"fmt"
 	"io"
-	"os"
-	"os/signal"
 	"strconv"
 	"strings"
 
@@ -45,14 +43,7 @@ func (b *ecsAPIService) Logs(ctx context.Context, options *cli.ProjectOptions, w
 		writer: writer,
 	}
 	err := b.SDK.GetLogs(ctx, name, consumer.Log)
-	if err != nil {
-		return err
-	}
-
-	signalChan := make(chan os.Signal, 1)
-	signal.Notify(signalChan, os.Interrupt)
-	<-signalChan
-	return nil
+	return err
 }
 
 func (l *logConsumer) Log(service, container, message string) {
diff --git a/ecs/sdk.go b/ecs/sdk.go
index 3f4429766..eb33d67ae 100644
--- a/ecs/sdk.go
+++ b/ecs/sdk.go
@@ -436,27 +436,32 @@ func (s sdk) GetLogs(ctx context.Context, name string, consumer func(service, co
 	logGroup := fmt.Sprintf("/docker-compose/%s", name)
 	var startTime int64
 	for {
-		var hasMore = true
-		var token *string
-		for hasMore {
-			events, err := s.CW.FilterLogEvents(&cloudwatchlogs.FilterLogEventsInput{
-				LogGroupName: aws.String(logGroup),
-				NextToken:    token,
-				StartTime:    aws.Int64(startTime),
-			})
-			if err != nil {
-				return err
-			}
-			if events.NextToken == nil {
-				hasMore = false
-			} else {
-				token = events.NextToken
-			}
+		select {
+		case <-ctx.Done():
+			return nil
+		default:
+			var hasMore = true
+			var token *string
+			for hasMore {
+				events, err := s.CW.FilterLogEvents(&cloudwatchlogs.FilterLogEventsInput{
+					LogGroupName: aws.String(logGroup),
+					NextToken:    token,
+					StartTime:    aws.Int64(startTime),
+				})
+				if err != nil {
+					return err
+				}
+				if events.NextToken == nil {
+					hasMore = false
+				} else {
+					token = events.NextToken
+				}
 
-			for _, event := range events.Events {
-				p := strings.Split(aws.StringValue(event.LogStreamName), "/")
-				consumer(p[1], p[2], aws.StringValue(event.Message))
-				startTime = *event.IngestionTime
+				for _, event := range events.Events {
+					p := strings.Split(aws.StringValue(event.LogStreamName), "/")
+					consumer(p[1], p[2], aws.StringValue(event.Message))
+					startTime = *event.IngestionTime
+				}
 			}
 		}
 		time.Sleep(500 * time.Millisecond)