From 2a5f2e10b06539adf172e866f4fdae6db936eb03 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Patryk=20Ma=C5=82ek?= <malekpatryk+github@gmail.com>
Date: Sat, 11 Nov 2023 15:35:09 +0100
Subject: [PATCH] fix: fix panic when getting container logs

---
 docker.go      | 29 ++++++++++++----
 docker_test.go | 90 ++++++++++++++++++++++++++++++++++++++++++++++++++
 2 files changed, 113 insertions(+), 6 deletions(-)

diff --git a/docker.go b/docker.go
index cb739a46d6..3a9bae551a 100644
--- a/docker.go
+++ b/docker.go
@@ -610,6 +610,11 @@ func (c *DockerContainer) CopyToContainer(ctx context.Context, fileContent []byt
 // StartLogProducer will start a concurrent process that will continuously read logs
 // from the container and will send them to each added LogConsumer
 func (c *DockerContainer) StartLogProducer(ctx context.Context) error {
+	// Short cirtuit.
+	if ctx.Err() != nil {
+		return ctx.Err()
+	}
+
 	if c.stopProducer != nil {
 		return errors.New("log producer already started")
 	}
@@ -636,14 +641,18 @@ func (c *DockerContainer) StartLogProducer(ctx context.Context) error {
 
 		r, err := c.provider.client.ContainerLogs(ctx, c.GetContainerID(), options)
 		if err != nil {
-			// if we can't get the logs, panic, we can't return an error to anything
-			// from within this goroutine
-			panic(err)
+			if errors.Is(err, context.Canceled) || errors.Is(err, context.DeadlineExceeded) {
+				return
+			}
+			// If the context is not canceled and deadline is not exceeded then retry.
+			goto BEGIN
 		}
 		defer c.provider.Close()
 
 		for {
 			select {
+			case <-ctx.Done():
+				return
 			case <-stop:
 				err := r.Close()
 				if err != nil {
@@ -714,9 +723,17 @@ func (c *DockerContainer) StartLogProducer(ctx context.Context) error {
 // and sending them to each added LogConsumer
 func (c *DockerContainer) StopLogProducer() error {
 	if c.stopProducer != nil {
-		c.stopProducer <- true
-		// block until the producer is actually done in order to avoid strange races
-		<-c.producerDone
+		// Producer might have been stopped already in case of an error.
+		// In that case, we don't want to block so we check if something is still
+		// listening on stopProducer or if we're able to read from the producerDone
+		// channel (this would indicate it's closed, hence producer goroutine has
+		// finished).
+		select {
+		case c.stopProducer <- true:
+		case <-c.producerDone:
+		default:
+		}
+
 		c.stopProducer = nil
 		c.producerDone = nil
 	}
diff --git a/docker_test.go b/docker_test.go
index 6d234cd4c7..fa803d57fe 100644
--- a/docker_test.go
+++ b/docker_test.go
@@ -32,6 +32,7 @@ const (
 	nginxDelayedImage = "docker.io/menedev/delayed-nginx:1.15.2"
 	nginxImage        = "docker.io/nginx"
 	nginxAlpineImage  = "docker.io/nginx:alpine"
+	pauseImage        = "registry.k8s.io/pause:3.9"
 	nginxDefaultPort  = "80/tcp"
 	nginxHighPort     = "8080/tcp"
 	daemonMaxVersion  = "1.41"
@@ -2159,3 +2160,92 @@ func TestImageBuiltFromDockerfile_KeepBuiltImage(t *testing.T) {
 		})
 	}
 }
+
+func TestLogs(t *testing.T) {
+	ctx := context.Background()
+
+	testcases := []struct {
+		name string
+		test func(ctx context.Context, t *testing.T, c Container)
+	}{
+		{
+			name: "StartLogProducer doesn't panic when provided context is cancelled",
+			test: func(ctx context.Context, t *testing.T, c Container) {
+				ctx, cancel := context.WithCancel(ctx)
+				cancel()
+
+				require.ErrorIs(t, c.StartLogProducer(ctx), context.Canceled)
+				t.Cleanup(func() { c.StopLogProducer() })
+			},
+		},
+		{
+			name: "StartLogProducer doesn't panic when provided context gets cancelled",
+			test: func(ctx context.Context, t *testing.T, c Container) {
+				ctx, cancel := context.WithTimeout(ctx, 100*time.Millisecond)
+				defer cancel()
+
+				require.NoError(t, c.StartLogProducer(ctx))
+				t.Cleanup(func() { c.StopLogProducer() })
+				<-ctx.Done()
+			},
+		},
+		{
+			name: "StartLogProducer and StopLogProducer work",
+			test: func(ctx context.Context, t *testing.T, c Container) {
+				ctx, cancel := context.WithCancel(ctx)
+				defer cancel()
+
+				require.NoError(t, c.StartLogProducer(ctx))
+				_, err := c.Logs(ctx)
+				require.NoError(t, err)
+				require.NoError(t, c.StopLogProducer())
+				_, err = c.Logs(ctx)
+				require.NoError(t, err)
+			},
+		},
+		{
+			name: "StartLogProducer doesn't panic when provided context exceeds deadline",
+			test: func(ctx context.Context, t *testing.T, c Container) {
+				ctx, cancel := context.WithDeadline(ctx, time.Now().Add(time.Nanosecond))
+				defer cancel()
+
+				require.ErrorIs(t, c.StartLogProducer(ctx), context.DeadlineExceeded)
+				t.Cleanup(func() { c.StopLogProducer() })
+			},
+		},
+		{
+			name: "Logs doesn't panic when context is cancelled",
+			test: func(ctx context.Context, t *testing.T, c Container) {
+				ctx, cancel := context.WithCancel(ctx)
+				cancel()
+
+				_, err := c.Logs(ctx)
+				require.ErrorIs(t, err, context.Canceled)
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		tc := tc
+
+		t.Run(tc.name, func(t *testing.T) {
+			t.Parallel()
+
+			container, err := GenericContainer(ctx,
+				GenericContainerRequest{
+					ContainerRequest: ContainerRequest{
+						Image: pauseImage,
+					},
+					Started: true,
+					Logger:  TestLogger(t),
+				})
+			require.NoError(t, err, container)
+
+			t.Cleanup(func() {
+				assert.NoError(t, container.Terminate(context.Background()))
+			})
+
+			tc.test(ctx, t, container)
+		})
+	}
+}