From 17b178f85a5046ac51b66be46fb87e3447763adc Mon Sep 17 00:00:00 2001
From: Steven Hartland <stevenmhartland@gmail.com>
Date: Fri, 16 Aug 2024 15:52:42 +0100
Subject: [PATCH] fix: logging restart (#2697)

Fix being able to stop and restart a container by moving the stop
processing from PreTerminates to PostStops hooks.

Fix log consumer being duplicated when restarting.

Fix race condition in log consumer test, cleaning up to flow to validate
correctly.
---
 lifecycle.go        | 19 ++++++------
 logconsumer_test.go | 74 +++++++++++++++++++++++++++++++--------------
 2 files changed, 61 insertions(+), 32 deletions(-)

diff --git a/lifecycle.go b/lifecycle.go
index 2a8ae949ab..40360a4c0b 100644
--- a/lifecycle.go
+++ b/lifecycle.go
@@ -165,26 +165,25 @@ var defaultCopyFileToContainerHook = func(files []ContainerFile) ContainerLifecy
 var defaultLogConsumersHook = func(cfg *LogConsumerConfig) ContainerLifecycleHooks {
 	return ContainerLifecycleHooks{
 		PostStarts: []ContainerHook{
-			// first post-start hook is to produce logs and start log consumers
+			// Produce logs sending details to the log consumers.
+			// See combineContainerHooks for the order of execution.
 			func(ctx context.Context, c Container) error {
-				dockerContainer := c.(*DockerContainer)
-
-				if cfg == nil {
+				if cfg == nil || len(cfg.Consumers) == 0 {
 					return nil
 				}
 
+				dockerContainer := c.(*DockerContainer)
+				dockerContainer.consumers = dockerContainer.consumers[:0]
 				for _, consumer := range cfg.Consumers {
 					dockerContainer.followOutput(consumer)
 				}
 
-				if len(cfg.Consumers) > 0 {
-					return dockerContainer.startLogProduction(ctx, cfg.Opts...)
-				}
-				return nil
+				return dockerContainer.startLogProduction(ctx, cfg.Opts...)
 			},
 		},
-		PreTerminates: []ContainerHook{
-			// first pre-terminate hook is to stop the log production
+		PostStops: []ContainerHook{
+			// Stop the log production.
+			// See combineContainerHooks for the order of execution.
 			func(ctx context.Context, c Container) error {
 				if cfg == nil || len(cfg.Consumers) == 0 {
 					return nil
diff --git a/logconsumer_test.go b/logconsumer_test.go
index 401e7a158f..6265f0a578 100644
--- a/logconsumer_test.go
+++ b/logconsumer_test.go
@@ -639,23 +639,54 @@ func Test_MultiContainerLogConsumer_CancelledContext(t *testing.T) {
 	assert.False(t, strings.Contains(actual, logStoppedForOutOfSyncMessage))
 }
 
+// FooLogConsumer is a test log consumer that accepts logs from the
+// "hello-world" Docker image, which prints out the "Hello from Docker!"
+// log message.
 type FooLogConsumer struct {
 	LogChannel chan string
+	t          *testing.T
 }
 
+// Accept receives a log message and sends it to the log channel if it
+// contains the "Hello from Docker!" message.
 func (c FooLogConsumer) Accept(rawLog Log) {
 	log := string(rawLog.Content)
-	c.LogChannel <- log
+	if strings.Contains(log, "Hello from Docker!") {
+		select {
+		case c.LogChannel <- log:
+		default:
+		}
+	}
+}
+
+// AssertRead waits for a log message to be received.
+func (c FooLogConsumer) AssertRead() {
+	select {
+	case <-c.LogChannel:
+	case <-time.After(5 * time.Second):
+		c.t.Fatal("receive timeout")
+	}
+}
+
+// SlurpOne reads a value from the channel if it is available.
+func (c FooLogConsumer) SlurpOne() {
+	select {
+	case <-c.LogChannel:
+	default:
+	}
 }
 
-func NewFooLogConsumer() *FooLogConsumer {
+func NewFooLogConsumer(t *testing.T) *FooLogConsumer {
+	t.Helper()
+
 	return &FooLogConsumer{
-		LogChannel: make(chan string),
+		t:          t,
+		LogChannel: make(chan string, 2),
 	}
 }
 
 func TestRestartContainerWithLogConsumer(t *testing.T) {
-	logConsumer := NewFooLogConsumer()
+	logConsumer := NewFooLogConsumer(t)
 
 	ctx := context.Background()
 	container, err := GenericContainer(ctx, GenericContainerRequest{
@@ -668,28 +699,27 @@ func TestRestartContainerWithLogConsumer(t *testing.T) {
 		},
 		Started: false,
 	})
-	if err != nil {
-		t.Fatalf("Cant create container: %s", err.Error())
-	}
+	terminateContainerOnEnd(t, ctx, container)
+	require.NoError(t, err)
 
+	// Start and confirm that the log consumer receives the log message.
 	err = container.Start(ctx)
-	if err != nil {
-		t.Fatalf("Cant start container: %s", err.Error())
-	}
+	require.NoError(t, err)
+
+	logConsumer.AssertRead()
 
-	d := 30 * time.Second
+	// Stop the container and clear any pending message.
+	d := 5 * time.Second
 	err = container.Stop(ctx, &d)
-	if err != nil {
-		t.Fatalf("Cant stop container: %s", err.Error())
-	}
+	require.NoError(t, err)
+
+	logConsumer.SlurpOne()
+
+	// Restart the container and confirm that the log consumer receives new log messages.
 	err = container.Start(ctx)
-	if err != nil {
-		t.Fatalf("Cant start container: %s", err.Error())
-	}
+	require.NoError(t, err)
 
-	for s := range logConsumer.LogChannel {
-		if strings.Contains(s, "Hello from Docker!") {
-			break
-		}
-	}
+	// First message is from the first start.
+	logConsumer.AssertRead()
+	logConsumer.AssertRead()
 }