From 1c3beb47a5c5286b00a6ea43cc1f987397c57158 Mon Sep 17 00:00:00 2001
From: HirazawaUi <695097494plus@gmail.com>
Date: Mon, 19 May 2025 22:03:21 +0800
Subject: [PATCH] Fix goroutine leak in unit tests

---
 pkg/kubelet/kubelet.go                       |  2 +-
 pkg/kubelet/watchdog/types.go                |  7 ++-
 pkg/kubelet/watchdog/watchdog_linux.go       |  7 +--
 pkg/kubelet/watchdog/watchdog_linux_test.go  | 55 ++++++--------------
 pkg/kubelet/watchdog/watchdog_unsupported.go |  7 ++-
 5 files changed, 32 insertions(+), 46 deletions(-)

diff --git a/pkg/kubelet/kubelet.go b/pkg/kubelet/kubelet.go
index 2a1a3e4aa2e..1e07229ebc3 100644
--- a/pkg/kubelet/kubelet.go
+++ b/pkg/kubelet/kubelet.go
@@ -1823,7 +1823,7 @@ func (kl *Kubelet) Run(updates <-chan kubetypes.PodUpdate) {
 	}
 
 	if utilfeature.DefaultFeatureGate.Enabled(features.SystemdWatchdog) {
-		kl.healthChecker.Start()
+		kl.healthChecker.Start(ctx)
 	}
 
 	kl.syncLoop(ctx, updates, kl)
diff --git a/pkg/kubelet/watchdog/types.go b/pkg/kubelet/watchdog/types.go
index 8f6d3838fde..a25bf2aeb38 100644
--- a/pkg/kubelet/watchdog/types.go
+++ b/pkg/kubelet/watchdog/types.go
@@ -16,11 +16,14 @@ limitations under the License.
 
 package watchdog
 
-import "net/http"
+import (
+	"context"
+	"net/http"
+)
 
 // HealthChecker defines the interface of health checkers.
 type HealthChecker interface {
-	Start()
+	Start(ctx context.Context)
 }
 
 // syncLoopHealthChecker contains the health check method for syncLoop.
diff --git a/pkg/kubelet/watchdog/watchdog_linux.go b/pkg/kubelet/watchdog/watchdog_linux.go
index e30275aaacd..abc4e0689f4 100644
--- a/pkg/kubelet/watchdog/watchdog_linux.go
+++ b/pkg/kubelet/watchdog/watchdog_linux.go
@@ -20,6 +20,7 @@ limitations under the License.
 package watchdog
 
 import (
+	"context"
 	"fmt"
 	"time"
 
@@ -122,20 +123,20 @@ func NewHealthChecker(syncLoop syncLoopHealthChecker, opts ...Option) (HealthChe
 	return hc, nil
 }
 
-func (hc *healthChecker) Start() {
+func (hc *healthChecker) Start(ctx context.Context) {
 	if hc.interval <= 0 {
 		klog.InfoS("Systemd watchdog is not enabled or the interval is invalid, so health checking will not be started.")
 		return
 	}
 	klog.InfoS("Starting systemd watchdog with interval", "interval", hc.interval)
 
-	go wait.Forever(func() {
+	go wait.UntilWithContext(ctx, func(ctx context.Context) {
 		if err := hc.doCheck(); err != nil {
 			klog.ErrorS(err, "Do not notify watchdog this iteration as the kubelet is reportedly not healthy")
 			return
 		}
 
-		err := wait.ExponentialBackoff(hc.retryBackoff, func() (bool, error) {
+		err := wait.ExponentialBackoffWithContext(ctx, hc.retryBackoff, func(_ context.Context) (bool, error) {
 			ack, err := hc.watchdog.SdNotify(false)
 			if err != nil {
 				klog.V(5).InfoS("Failed to notify systemd watchdog, retrying", "error", err)
diff --git a/pkg/kubelet/watchdog/watchdog_linux_test.go b/pkg/kubelet/watchdog/watchdog_linux_test.go
index ed167742523..80d8f286ebf 100644
--- a/pkg/kubelet/watchdog/watchdog_linux_test.go
+++ b/pkg/kubelet/watchdog/watchdog_linux_test.go
@@ -23,13 +23,14 @@ import (
 	"bytes"
 	"errors"
 	"flag"
+	"io"
 	"net/http"
 	"strings"
-	"sync"
 	"testing"
 	"time"
 
 	"k8s.io/klog/v2"
+	"k8s.io/kubernetes/test/utils/ktesting"
 )
 
 // Mock syncLoopHealthChecker
@@ -139,8 +140,19 @@ func TestHealthCheckerStart(t *testing.T) {
 
 	for _, tt := range tests {
 		t.Run(tt.name, func(t *testing.T) {
+			tCtx := ktesting.Init(t)
+			defer func() {
+				tCtx.Cancel("test has completed")
+			}()
 			// Capture logs
-			logBuffer := setupLogging(t)
+			var logBuffer bytes.Buffer
+			flags := &flag.FlagSet{}
+			klog.InitFlags(flags)
+			if err := flags.Set("v", "5"); err != nil {
+				t.Fatal(err)
+			}
+			klog.LogToStderr(false)
+			klog.SetOutput(&logBuffer)
 
 			// Mock SdWatchdogEnabled to return a valid value
 			mockClient := &mockWatchdogClient{
@@ -156,13 +168,15 @@ func TestHealthCheckerStart(t *testing.T) {
 			}
 
 			// Start the health checker
-			hc.Start()
+			hc.Start(tCtx)
 
 			// Wait for a short period to allow the health check to run
 			time.Sleep(2 * interval)
 
 			// Check logs to verify the health check ran
 			klog.Flush()
+			// prevent further writes into buf
+			klog.SetOutput(io.Discard)
 			logs := logBuffer.String()
 			for _, expectedLog := range tt.expectedLogs {
 				if !strings.Contains(logs, expectedLog) {
@@ -172,38 +186,3 @@ func TestHealthCheckerStart(t *testing.T) {
 		})
 	}
 }
-
-// threadSafeBuffer is a thread-safe wrapper around bytes.Buffer.
-type threadSafeBuffer struct {
-	buffer bytes.Buffer
-	mu     sync.Mutex
-}
-
-func (b *threadSafeBuffer) Write(p []byte) (n int, err error) {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-	return b.buffer.Write(p)
-}
-
-func (b *threadSafeBuffer) String() string {
-	b.mu.Lock()
-	defer b.mu.Unlock()
-	return b.buffer.String()
-}
-
-// setupLogging sets up logging to capture output using a thread-safe buffer.
-func setupLogging(t *testing.T) *threadSafeBuffer {
-	flags := &flag.FlagSet{}
-	klog.InitFlags(flags)
-	if err := flags.Set("v", "5"); err != nil {
-		t.Fatal(err)
-	}
-	klog.LogToStderr(false)
-
-	logBuffer := &threadSafeBuffer{}
-
-	// Set the output to the thread-safe buffer
-	klog.SetOutput(logBuffer)
-
-	return logBuffer
-}
diff --git a/pkg/kubelet/watchdog/watchdog_unsupported.go b/pkg/kubelet/watchdog/watchdog_unsupported.go
index b536af21f2d..be54f932ec3 100644
--- a/pkg/kubelet/watchdog/watchdog_unsupported.go
+++ b/pkg/kubelet/watchdog/watchdog_unsupported.go
@@ -19,7 +19,10 @@ limitations under the License.
 
 package watchdog
 
-import "k8s.io/apiserver/pkg/server/healthz"
+import (
+	"context"
+	"k8s.io/apiserver/pkg/server/healthz"
+)
 
 type healthCheckerUnsupported struct{}
 
@@ -36,6 +39,6 @@ func NewHealthChecker(_ syncLoopHealthChecker, _ ...Option) (HealthChecker, erro
 	return &healthCheckerUnsupported{}, nil
 }
 
-func (ow *healthCheckerUnsupported) Start() {
+func (ow *healthCheckerUnsupported) Start(ctx context.Context) {
 	return
 }