From c3e796510e90ad9b04b751ffce0f79670f582592 Mon Sep 17 00:00:00 2001
From: Curtis Robert <crobert@splunk.com>
Date: Fri, 5 Apr 2024 07:28:39 -0700
Subject: [PATCH] [chore][receiver/zookeeper] Enable goleak check (#32179)

**Description:** <Describe what has changed.>
<!--Ex. Fixing a bug - Describe the bug and how this fixes the issue.
Ex. Adding a feature - Explain what this achieves.-->
This change enables `goleak` to help ensure the ZooKeeper receiver does
not leak any goroutines. This is a test only change, the existing test
implementation needed a bit of modification to properly close the
network listener.

**Link to tracking Issue:** <Issue number if applicable>
#30438

**Testing:** <Describe what testing was performed and which tests were
added.>
Existing tests and new `goleak` check are all passing.
---
 receiver/zookeeperreceiver/package_test.go | 14 ++++++++
 receiver/zookeeperreceiver/scraper_test.go | 37 ++++++++++++++++------
 2 files changed, 42 insertions(+), 9 deletions(-)
 create mode 100644 receiver/zookeeperreceiver/package_test.go

diff --git a/receiver/zookeeperreceiver/package_test.go b/receiver/zookeeperreceiver/package_test.go
new file mode 100644
index 000000000000..77394b47128a
--- /dev/null
+++ b/receiver/zookeeperreceiver/package_test.go
@@ -0,0 +1,14 @@
+// Copyright The OpenTelemetry Authors
+// SPDX-License-Identifier: Apache-2.0
+
+package zookeeperreceiver
+
+import (
+	"testing"
+
+	"go.uber.org/goleak"
+)
+
+func TestMain(m *testing.M) {
+	goleak.VerifyTestMain(m)
+}
diff --git a/receiver/zookeeperreceiver/scraper_test.go b/receiver/zookeeperreceiver/scraper_test.go
index d935b669725a..92309f0566dd 100644
--- a/receiver/zookeeperreceiver/scraper_test.go
+++ b/receiver/zookeeperreceiver/scraper_test.go
@@ -270,8 +270,16 @@ func TestZookeeperMetricsScraperScrape(t *testing.T) {
 		t.Run(tt.name, func(t *testing.T) {
 			localAddr := testutil.GetAvailableLocalAddress(t)
 			if !tt.mockZKConnectionErr {
-				ms := mockedServer{ready: make(chan bool, 1)}
-				go ms.mockZKServer(t, localAddr, tt.mockedZKCmdToOutputFilename)
+				listener, err := net.Listen("tcp", localAddr)
+				require.NoError(t, err)
+				ms := mockedServer{
+					listener: listener,
+					ready:    make(chan bool, 1),
+					quit:     make(chan struct{}),
+				}
+
+				defer ms.shutdown()
+				go ms.mockZKServer(t, tt.mockedZKCmdToOutputFilename)
 				<-ms.ready
 			}
 
@@ -337,19 +345,26 @@ func TestZookeeperShutdownBeforeScrape(t *testing.T) {
 }
 
 type mockedServer struct {
+	listener net.Listener
+
 	ready chan bool
+	quit  chan struct{}
 }
 
-func (ms *mockedServer) mockZKServer(t *testing.T, endpoint string, cmdToFileMap map[string]string) {
+func (ms *mockedServer) mockZKServer(t *testing.T, cmdToFileMap map[string]string) {
 	var cmd string
-	listener, err := net.Listen("tcp", endpoint)
-	require.NoError(t, err)
-	defer listener.Close()
 	ms.ready <- true
 
 	for {
-		conn, err := listener.Accept()
-		require.NoError(t, err)
+		conn, err := ms.listener.Accept()
+		if err != nil {
+			select {
+			case <-ms.quit:
+				return
+			default:
+				require.NoError(t, err)
+			}
+		}
 		reader := bufio.NewReader(conn)
 		scanner := bufio.NewScanner(reader)
 		scanner.Scan()
@@ -366,6 +381,10 @@ func (ms *mockedServer) mockZKServer(t *testing.T, endpoint string, cmdToFileMap
 		require.NoError(t, err)
 
 		conn.Close()
-
 	}
 }
+
+func (ms *mockedServer) shutdown() {
+	close(ms.quit)
+	ms.listener.Close()
+}