Add go leak checks to ES/OS tests (jaegertracing#6339)

## Which problem is this PR solving?
- Fixes a part of jaegertracing#5006 

## Description of the changes
- This is a better version of `elasticsearch_test` reducing the number
of leaks from 50 to less than 10. Unfortunately we can't fix all the
leakages because there is no way to close the v8 client of ES leading to
some unclosed http transports. If the changes looks good then we can
merge without embedding `goleak` (as it will lead to failure of test)

## How was this change tested?
- By running the integration tests

## Checklist
- [x] I have read
https://github.com/jaegertracing/jaeger/blob/master/CONTRIBUTING_GUIDELINES.md
- [x] I have signed all commits
- [x] I have added unit tests for the new functionality
- [ ] I have run lint and test steps successfully
  - for `jaeger`: `make lint test`
  - for `jaeger-ui`: `yarn lint` and `yarn test`

---------

Signed-off-by: Manik2708 <[email protected]>

pkg/testutils/leakcheck.go

@@ -34,10 +34,31 @@ func IgnoreGoMetricsMeterLeak() goleak.Option {
 	return goleak.IgnoreTopFunction("github.com/rcrowley/go-metrics.(*meterArbiter).tick")
 }
 
+// Don't use this in any other method other than leaks for ElasticSearch and OpenSearch
+// These leaks are from olivere client not from the jaeger
+// See this PR for context: https://github.com/jaegertracing/jaeger/pull/6339
+func ignoreHttpTransportWriteLoopLeak() goleak.Option {
+	return goleak.IgnoreTopFunction("net/http.(*persistConn).writeLoop")
+}
+
+// Don't use this in any other method other than leaks for ElasticSearch and OpenSearch
+// These leaks are from olivere client not from the jaeger
+// See this PR for context: https://github.com/jaegertracing/jaeger/pull/6339
+func ignoreHttpTransportPollRuntimeLeak() goleak.Option {
+	return goleak.IgnoreTopFunction("internal/poll.runtime_pollWait")
+}
+
+// Don't use this in any other method other than leaks for ElasticSearch and OpenSearch
+// These leaks are from olivere client not from the jaeger
+// See this PR for context: https://github.com/jaegertracing/jaeger/pull/6339
+func ignoreHttpTransportReadLoopLeak() goleak.Option {
+	return goleak.IgnoreTopFunction("net/http.(*persistConn).readLoop")
+}
+
 // VerifyGoLeaks verifies that unit tests do not leak any goroutines.
 // It should be called in TestMain.
 func VerifyGoLeaks(m *testing.M) {
-	goleak.VerifyTestMain(m, IgnoreGlogFlushDaemonLeak(), IgnoreOpenCensusWorkerLeak())
+	goleak.VerifyTestMain(m, IgnoreGlogFlushDaemonLeak(), IgnoreOpenCensusWorkerLeak(), IgnoreGoMetricsMeterLeak())
 }
 
 // VerifyGoLeaksOnce verifies that a given unit test does not leak any goroutines.
@@ -49,3 +70,15 @@ func VerifyGoLeaks(m *testing.M) {
 func VerifyGoLeaksOnce(t *testing.T) {
 	goleak.VerifyNone(t, IgnoreGlogFlushDaemonLeak(), IgnoreOpenCensusWorkerLeak(), IgnoreGoMetricsMeterLeak())
 }
+
+// VerifyGoLeaksOnceForES is go leak check for ElasticSearch integration tests (v1)
+// This must not be used anywhere else other than integration package for v1
+func VerifyGoLeaksOnceForES(t *testing.T) {
+	goleak.VerifyNone(t, ignoreHttpTransportWriteLoopLeak(), ignoreHttpTransportPollRuntimeLeak(), ignoreHttpTransportReadLoopLeak())
+}
+
+// VerifyGoLeaksForES is go leak check for integration package in ElasticSearch Environment
+// This must not be used anywhere else other than integration package in ES environment for v1
+func VerifyGoLeaksForES(m *testing.M) {
+	goleak.VerifyTestMain(m, ignoreHttpTransportWriteLoopLeak(), ignoreHttpTransportPollRuntimeLeak(), ignoreHttpTransportReadLoopLeak())
+}

plugin/storage/integration/elasticsearch_test.go

@@ -23,6 +23,7 @@ import (
 
 	"github.com/jaegertracing/jaeger/pkg/config"
 	"github.com/jaegertracing/jaeger/pkg/metrics"
+	"github.com/jaegertracing/jaeger/pkg/testutils"
 	"github.com/jaegertracing/jaeger/plugin/storage/es"
 	"github.com/jaegertracing/jaeger/storage/dependencystore"
 )
@@ -71,16 +72,20 @@ func (s *ESStorageIntegration) getVersion() (uint, error) {
 	return uint(esVersion), nil
 }
 
-func (s *ESStorageIntegration) initializeES(t *testing.T, allTagsAsFields bool) {
+func (s *ESStorageIntegration) initializeES(t *testing.T, c *http.Client, allTagsAsFields bool) {
 	rawClient, err := elastic.NewClient(
 		elastic.SetURL(queryURL),
-		elastic.SetSniff(false))
+		elastic.SetSniff(false),
+		elastic.SetHttpClient(c))
 	require.NoError(t, err)
-
+	t.Cleanup(func() {
+		rawClient.Stop()
+	})
 	s.client = rawClient
 	s.v8Client, err = elasticsearch8.NewClient(elasticsearch8.Config{
 		Addresses:            []string{queryURL},
 		DiscoverNodesOnStart: false,
+		Transport:            c.Transport,
 	})
 	require.NoError(t, err)
 
@@ -144,10 +149,10 @@ func (s *ESStorageIntegration) initSpanstore(t *testing.T, allTagsAsFields bool)
 	require.NoError(t, err)
 }
 
-func healthCheck() error {
+func healthCheck(c *http.Client) error {
 	for i := 0; i < 200; i++ {
-		if _, err := http.Get(queryURL); err == nil {
-			return nil
+		if resp, err := c.Get(queryURL); err == nil {
+			return resp.Body.Close()
 		}
 		time.Sleep(100 * time.Millisecond)
 	}
@@ -156,7 +161,8 @@ func healthCheck() error {
 
 func testElasticsearchStorage(t *testing.T, allTagsAsFields bool) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
-	require.NoError(t, healthCheck())
+	c := getESHttpClient(t)
+	require.NoError(t, healthCheck(c))
 	s := &ESStorageIntegration{
 		StorageIntegration: StorageIntegration{
 			Fixtures:        LoadAndParseQueryTestCases(t, "fixtures/queries_es.json"),
@@ -166,23 +172,33 @@ func testElasticsearchStorage(t *testing.T, allTagsAsFields bool) {
 			GetOperationsMissingSpanKind: true,
 		},
 	}
-	s.initializeES(t, allTagsAsFields)
+	s.initializeES(t, c, allTagsAsFields)
 	s.RunAll(t)
 }
 
 func TestElasticsearchStorage(t *testing.T) {
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
 	testElasticsearchStorage(t, false)
 }
 
 func TestElasticsearchStorage_AllTagsAsObjectFields(t *testing.T) {
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
 	testElasticsearchStorage(t, true)
 }
 
 func TestElasticsearchStorage_IndexTemplates(t *testing.T) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
-	require.NoError(t, healthCheck())
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
+	c := getESHttpClient(t)
+	require.NoError(t, healthCheck(c))
 	s := &ESStorageIntegration{}
-	s.initializeES(t, true)
+	s.initializeES(t, c, true)
 	esVersion, err := s.getVersion()
 	require.NoError(t, err)
 	// TODO abstract this into pkg/es/client.IndexManagementLifecycleAPI

plugin/storage/integration/es_index_cleaner_test.go

@@ -6,13 +6,16 @@ package integration
 import (
 	"context"
 	"fmt"
+	"net/http"
 	"os/exec"
 	"testing"
 
 	elasticsearch8 "github.com/elastic/go-elasticsearch/v8"
 	"github.com/olivere/elastic"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
 )
 
 const (
@@ -28,7 +31,10 @@ const (
 
 func TestIndexCleaner_doNotFailOnEmptyStorage(t *testing.T) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
-	client, err := createESClient()
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
+	client, err := createESClient(t, getESHttpClient(t))
 	require.NoError(t, err)
 	_, err = client.DeleteIndex("*").Do(context.Background())
 	require.NoError(t, err)
@@ -48,7 +54,10 @@ func TestIndexCleaner_doNotFailOnEmptyStorage(t *testing.T) {
 
 func TestIndexCleaner_doNotFailOnFullStorage(t *testing.T) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
-	client, err := createESClient()
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
+	client, err := createESClient(t, getESHttpClient(t))
 	require.NoError(t, err)
 	tests := []struct {
 		envs []string
@@ -70,9 +79,13 @@ func TestIndexCleaner_doNotFailOnFullStorage(t *testing.T) {
 
 func TestIndexCleaner(t *testing.T) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
-	client, err := createESClient()
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
+	hcl := getESHttpClient(t)
+	client, err := createESClient(t, hcl)
 	require.NoError(t, err)
-	v8Client, err := createESV8Client()
+	v8Client, err := createESV8Client(hcl.Transport)
 	require.NoError(t, err)
 
 	tests := []struct {
@@ -223,16 +236,24 @@ func runEsRollover(action string, envs []string, adaptiveSampling bool) error {
 	return err
 }
 
-func createESClient() (*elastic.Client, error) {
-	return elastic.NewClient(
+func createESClient(t *testing.T, hcl *http.Client) (*elastic.Client, error) {
+	cl, err := elastic.NewClient(
 		elastic.SetURL(queryURL),
-		elastic.SetSniff(false))
+		elastic.SetSniff(false),
+		elastic.SetHttpClient(hcl),
+	)
+	require.NoError(t, err)
+	t.Cleanup(func() {
+		cl.Stop()
+	})
+	return cl, nil
 }
 
-func createESV8Client() (*elasticsearch8.Client, error) {
+func createESV8Client(tr http.RoundTripper) (*elasticsearch8.Client, error) {
 	return elasticsearch8.NewClient(elasticsearch8.Config{
 		Addresses:            []string{queryURL},
 		DiscoverNodesOnStart: false,
+		Transport:            tr,
 	})
 }
 
@@ -243,3 +264,11 @@ func cleanESIndexTemplates(t *testing.T, client *elastic.Client, v8Client *elast
 	}
 	s.cleanESIndexTemplates(t, prefix)
 }
+
+func getESHttpClient(t *testing.T) *http.Client {
+	tr := &http.Transport{}
+	t.Cleanup(func() {
+		tr.CloseIdleConnections()
+	})
+	return &http.Client{Transport: tr}
+}

plugin/storage/integration/es_index_rollover_test.go

@@ -11,6 +11,8 @@ import (
 	"github.com/olivere/elastic"
 	"github.com/stretchr/testify/assert"
 	"github.com/stretchr/testify/require"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
 )
 
 const (
@@ -19,7 +21,10 @@ const (
 
 func TestIndexRollover_FailIfILMNotPresent(t *testing.T) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
-	client, err := createESClient()
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
+	client, err := createESClient(t, getESHttpClient(t))
 	require.NoError(t, err)
 	require.NoError(t, err)
 	// make sure ES is clean
@@ -35,6 +40,9 @@ func TestIndexRollover_FailIfILMNotPresent(t *testing.T) {
 
 func TestIndexRollover_CreateIndicesWithILM(t *testing.T) {
 	SkipUnlessEnv(t, "elasticsearch", "opensearch")
+	t.Cleanup(func() {
+		testutils.VerifyGoLeaksOnceForES(t)
+	})
 	// Test using the default ILM Policy Name, i.e. do not pass the ES_ILM_POLICY_NAME env var to the rollover script.
 	t.Run("DefaultPolicyName", func(t *testing.T) {
 		runCreateIndicesWithILM(t, defaultILMPolicyName)
@@ -47,7 +55,7 @@ func TestIndexRollover_CreateIndicesWithILM(t *testing.T) {
 }
 
 func runCreateIndicesWithILM(t *testing.T, ilmPolicyName string) {
-	client, err := createESClient()
+	client, err := createESClient(t, getESHttpClient(t))
 	require.NoError(t, err)
 	esVersion, err := getVersion(client)
 	require.NoError(t, err)
@@ -82,7 +90,7 @@ func runIndexRolloverWithILMTest(t *testing.T, client *elastic.Client, prefix st
 	}
 	// make sure ES is cleaned before test
 	cleanES(t, client, ilmPolicyName)
-	v8Client, err := createESV8Client()
+	v8Client, err := createESV8Client(getESHttpClient(t).Transport)
 	require.NoError(t, err)
 	// make sure ES is cleaned after test
 	defer cleanES(t, client, ilmPolicyName)

plugin/storage/integration/package_test.go

@@ -0,0 +1,19 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package integration
+
+import (
+	"os"
+	"testing"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
+)
+
+func TestMain(m *testing.M) {
+	if os.Getenv("STORAGE") == "elasticsearch" || os.Getenv("STORAGE") == "opensearch" {
+		testutils.VerifyGoLeaksForES(m)
+	} else {
+		testutils.VerifyGoLeaks(m)
+	}
+}

scripts/check-goleak-files.sh

@@ -8,6 +8,7 @@ set -euo pipefail
 bad_pkgs=0
 total_pkgs=0
 failed_pkgs=0
+invalid_use_pkgs=0
 
 # shellcheck disable=SC2048
 for dir in $*; do
@@ -20,18 +21,25 @@ for dir in $*; do
     continue
   fi
   good=0
+  invalid=0
   for test in ${testFiles}; do
     if grep -q "TestMain" "${test}" && grep -q "testutils.VerifyGoLeaks" "${test}"; then
+      if [ "${dir}" != "./plugin/storage/integration/" ] &&  grep -q "testutils.VerifyGoLeaksForES" "${test}"; then
+          invalid=1
+          break
+      fi
       good=1
       break
     fi
   done
+
   if ((good == 0)); then
-    echo "Error(check-goleak): no goleak check in package ${dir}"
-    ((bad_pkgs+=1))
-    if [[ "${dir}" == "./cmd/jaeger/internal/integration/" || "${dir}" == "./plugin/storage/integration/" ]]; then
-      echo "	... this package is temporarily allowed and will not cause linter failure"
+    if ((invalid == 1)); then
+      echo "Error(check-goleak): VerifyGoLeaksForES should only be used in integration package but it is used in ${dir} also"
+      ((invalid_use_pkgs+=1))
     else
+      echo "Error(check-goleak): no goleak check in package ${dir}"
+      ((bad_pkgs+=1))
       ((failed_pkgs+=1))
     fi
   fi
@@ -45,6 +53,10 @@ if ((failed_pkgs > 0)); then
   echo "⛔ Fatal(check-goleak): no goleak check in ${bad_pkgs} package(s), ${failed_pkgs} of which not allowed."
   help
   exit 1
+elif ((invalid_use_pkgs > 0)); then
+  echo "⛔ Fatal(check-goleak): use of VerifyGoLeaksForES in package(s) ${invalid_use_pkgs} which is not allowed"
+  help
+  exit 1
 elif ((bad_pkgs > 0)); then
   echo "🐞 Warning(check-goleak): no goleak check in ${bad_pkgs} package(s)."
   help