From be99d25fb259c6cb036390cffc8645c1c0b0609b Mon Sep 17 00:00:00 2001
From: Rob Watson <rob@netflux.io>
Date: Sat, 11 Jun 2022 07:13:42 +0200
Subject: [PATCH] podWatcher: prefer channel for synchronization

---
 logs/pod_watcher.go      | 21 ++++++++++++---------
 logs/pod_watcher_test.go |  2 --
 logs/watcher.go          | 20 ++------------------
 3 files changed, 14 insertions(+), 29 deletions(-)

diff --git a/logs/pod_watcher.go b/logs/pod_watcher.go
index ae0b678..4bd6310 100644
--- a/logs/pod_watcher.go
+++ b/logs/pod_watcher.go
@@ -101,10 +101,10 @@ func (pw *PodWatcher) watchPods(ctx context.Context, wg *sync.WaitGroup) error {
 	ticker := time.NewTicker(tickerInterval)
 	defer ticker.Stop()
 
-	// streamErrors is never closed.
+	logStream := make(chan string)
 	streamErrors := make(chan *streamError)
-	resultChan := watcher.ResultChan()
 
+	resultChan := watcher.ResultChan()
 	for {
 		select {
 		case <-ctx.Done():
@@ -123,7 +123,7 @@ func (pw *PodWatcher) watchPods(ctx context.Context, wg *sync.WaitGroup) error {
 					pw.status[pod.Name] = true
 					wg.Add(1)
 					go func() {
-						if err := copyPodLogs(ctx, wg, pw.client, pod, pw.container, pw.dst); err != nil {
+						if err := copyPodLogs(ctx, wg, pw.client, pod, pw.container, logStream); err != nil {
 							streamErrors <- err
 						}
 					}()
@@ -137,6 +137,13 @@ func (pw *PodWatcher) watchPods(ctx context.Context, wg *sync.WaitGroup) error {
 				}
 			}
 
+		// logStream is never closed
+		case l := <-logStream:
+			if _, err := pw.dst.Write([]byte(l)); err != nil {
+				return fmt.Errorf("error writing logs: %v", err)
+			}
+
+		// streamErrors is never closed
 		case streamErr := <-streamErrors:
 			if streamErr.recoverable {
 				// if the error is recoverable, we just remove the pod from the status
@@ -169,19 +176,17 @@ func (pw *PodWatcher) removePod(podName string) {
 	delete(pw.status, podName)
 }
 
-func copyPodLogs(ctx context.Context, wg *sync.WaitGroup, client KubernetesClient, pod *corev1.Pod, container string, dst io.Writer) *streamError {
+func copyPodLogs(ctx context.Context, wg *sync.WaitGroup, client KubernetesClient, pod *corev1.Pod, container string, logStream chan string) *streamError {
 	defer wg.Done()
 
 	podLogOpts := corev1.PodLogOptions{
 		Follow:    true,
 		Container: container,
 	}
-
 	req := client.Typed.CoreV1().Pods(pod.Namespace).GetLogs(pod.Name, &podLogOpts)
 	logs, err := req.Stream(ctx)
 
 	// If one container is still being created, do not treat this as a fatal error.
-	// We try to verify the error type as strictly as possible.
 	var statusErr *apierrors.StatusError
 	if errors.As(err, &statusErr) && statusErr.Status().Reason == metav1.StatusReasonBadRequest && strings.Contains(statusErr.Error(), "ContainerCreating") {
 		return newRecoverableError(err, pod.Name)
@@ -193,9 +198,7 @@ func copyPodLogs(ctx context.Context, wg *sync.WaitGroup, client KubernetesClien
 
 	scanner := bufio.NewScanner(logs)
 	for scanner.Scan() {
-		if _, err = dst.Write([]byte("[" + pod.Name + "] " + scanner.Text() + nl)); err != nil {
-			return newStreamError(fmt.Errorf("error writing logs: %v", err), pod.Name)
-		}
+		logStream <- "[" + pod.Name + "] " + scanner.Text() + nl
 	}
 	if err := scanner.Err(); err != nil {
 		return newStreamError(fmt.Errorf("error scanning logs: %v", err), pod.Name)
diff --git a/logs/pod_watcher_test.go b/logs/pod_watcher_test.go
index df3a0b5..87e117a 100644
--- a/logs/pod_watcher_test.go
+++ b/logs/pod_watcher_test.go
@@ -6,7 +6,6 @@ import (
 	"errors"
 	"fmt"
 	"io"
-	"log"
 	"net/http"
 	"strings"
 	"testing"
@@ -134,7 +133,6 @@ func TestPodWatcher(t *testing.T) {
 	for _, tc := range testCases {
 		t.Run(tc.name, func(t *testing.T) {
 			podsWatcher := watch.NewFake()
-			defer log.Println("exiting test func")
 
 			clientset := mockClientset{
 				getLogsRespBody:   tc.getLogsRespBody,
diff --git a/logs/watcher.go b/logs/watcher.go
index e2f31b5..b4f958d 100644
--- a/logs/watcher.go
+++ b/logs/watcher.go
@@ -7,7 +7,6 @@ import (
 	"io"
 	"log"
 	"os"
-	"sync"
 	"time"
 
 	corev1 "k8s.io/api/core/v1"
@@ -29,21 +28,6 @@ type KubernetesClient struct {
 	Untyped dynamic.Interface
 }
 
-// concurrentWriter implements an io.Writer that can be safely written to from
-// multiple goroutines.
-type concurrentWriter struct {
-	w  io.Writer
-	mu sync.Mutex
-}
-
-// Write implements io.Writer.
-func (cw *concurrentWriter) Write(p []byte) (int, error) {
-	cw.mu.Lock()
-	defer cw.mu.Unlock()
-
-	return cw.w.Write(p)
-}
-
 type PodWatcherInterface interface {
 	WatchPods(ctx context.Context) error
 	Close()
@@ -71,7 +55,7 @@ type Watcher struct {
 	podWatcher     PodWatcherInterface
 	podWatcherFunc PodWatcherFunc
 	errChan        chan error
-	dst            *concurrentWriter
+	dst            io.Writer
 	logger         *log.Logger
 }
 
@@ -82,7 +66,7 @@ func NewWatcher(params WatcherParams, client KubernetesClient, podWatcherFunc Po
 		client:         client,
 		podWatcherFunc: podWatcherFunc,
 		errChan:        make(chan error),
-		dst:            &concurrentWriter{w: dst},
+		dst:            dst,
 		logger:         logger,
 	}
 }