Refactor logstream

Author: lukaszcl

lib/logstream/logstream.go

@@ -441,15 +441,34 @@ func (m *LogStream) SaveLogTargetsLocations(writer LogWriter) {
 
 // Stop stops the consumer and closes listening channel (it won't be accepting any logs from now on)
 func (g *ContainerLogConsumer) stop() error {
-	if g.isDone {
-		return nil
-	}
+	var stopErr error
+
+	// Use sync.Once to ensure the stop logic is executed only once
+	g.stopOnce.Do(func() {
+		// Send a signal to the listening channel if possible
+		select {
+		case g.logListeningDone <- struct{}{}:
+		default:
+			// No action needed if the channel is already full
+		}
 
-	g.isDone = true
-	g.logListeningDone <- struct{}{}
-	defer close(g.logListeningDone)
+		// Close the channel
+		close(g.logListeningDone)
+	})
 
-	return nil
+	return stopErr
+}
+
+// IsStopped checks if the consumer has already been stopped
+func (g *ContainerLogConsumer) IsStopped() bool {
+	select {
+	case <-g.logListeningDone:
+		// Channel has been closed, meaning stop() was called
+		return true
+	default:
+		// Channel is still open
+		return false
+	}
 }
 
 // DisconnectContainer disconnects particular container
@@ -558,7 +577,7 @@ func (m *LogStream) GetAllLogsAndConsume(preExecuteFn ConsumerConsumingFn, consu
 		// set the cursor to the end of the file, when done to resume writing, unless it was closed
 		//revive:disable
 		defer func() {
-			if !consumer.isDone {
+			if !consumer.IsStopped() {
 				_, deferErr := consumer.tempFile.Seek(0, 2)
 				attachError(deferErr)
 			}
@@ -623,8 +642,10 @@ func (m *LogStream) GetAllLogsAndConsume(preExecuteFn ConsumerConsumingFn, consu
 // FlushLogsToTargets flushes all logs for all consumers (containers) to their targets
 func (m *LogStream) FlushLogsToTargets() error {
 	var preExecuteFn = func(consumer *ContainerLogConsumer) error {
-		// do not accept any new logs
-		consumer.isDone = true
+		// Stop the consumer to ensure it doesn't accept new logs
+		if err := consumer.stop(); err != nil {
+			return fmt.Errorf("failed to stop consumer %s: %w", consumer.name, err)
+		}
 
 		for _, handler := range m.logTargetHandlers {
 			consumer.ls.log.Debug().
@@ -699,7 +720,7 @@ type ContainerLogConsumer struct {
 	ls               *LogStream
 	tempFile         *os.File
 	encoder          *gob.Encoder
-	isDone           bool
+	stopOnce         sync.Once
 	hasErrored       bool
 	logListeningDone chan struct{}
 	container        LogProducingContainer
@@ -720,7 +741,6 @@ func newContainerLogConsumer(ctx context.Context, lw *LogStream, container LogPr
 		prefix:           prefix,
 		logTargets:       logTargets,
 		ls:               lw,
-		isDone:           false,
 		hasErrored:       false,
 		logListeningDone: make(chan struct{}, 1),
 		container:        container,
@@ -769,8 +789,7 @@ func (g *ContainerLogConsumer) ResetTempFile() error {
 // MarkAsErrored marks the consumer as errored (which makes it stop accepting logs)
 func (g *ContainerLogConsumer) MarkAsErrored() {
 	g.hasErrored = true
-	g.isDone = true
-	close(g.logListeningDone)
+	_ = g.stop()
 }
 
 // GetContainer returns the container that this consumer is connected to
@@ -787,7 +806,8 @@ func (g *ContainerLogConsumer) Accept(l tc.Log) {
 		return
 	}
 
-	if g.isDone {
+	if g.IsStopped() {
+		// if the consumer is stopped, we don't want to accept any more logs
 		return
 	}