Fix groutine leaks in cmd/anonymizer/app/writer (jaegertracing#5216)

cmd/anonymizer/app/anonymizer/anonymizer.go

@@ -15,6 +15,7 @@
 package anonymizer
 
 import (
+	"context"
 	"encoding/json"
 	"fmt"
 	"hash/fnv"
@@ -57,6 +58,8 @@ type Anonymizer struct {
 	lock        sync.Mutex
 	mapping     mapping
 	options     Options
+	cancel      context.CancelFunc
+	wg          sync.WaitGroup
 }
 
 // Options represents the various options with which the anonymizer can be configured.
@@ -70,6 +73,7 @@ type Options struct {
 // New creates new Anonymizer. The mappingFile stores the mapping from original to
 // obfuscated strings, in case later investigations require looking at the original traces.
 func New(mappingFile string, options Options, logger *zap.Logger) *Anonymizer {
+	ctx, cancel := context.WithCancel(context.Background())
 	a := &Anonymizer{
 		mappingFile: mappingFile,
 		logger:      logger,
@@ -78,6 +82,7 @@ func New(mappingFile string, options Options, logger *zap.Logger) *Anonymizer {
 			Operations: make(map[string]string),
 		},
 		options: options,
+		cancel:  cancel,
 	}
 	if _, err := os.Stat(filepath.Clean(mappingFile)); err == nil {
 		dat, err := os.ReadFile(filepath.Clean(mappingFile))
@@ -88,14 +93,28 @@ func New(mappingFile string, options Options, logger *zap.Logger) *Anonymizer {
 			logger.Fatal("Cannot unmarshal previous mapping", zap.Error(err))
 		}
 	}
+	a.wg.Add(1)
 	go func() {
-		for range time.NewTicker(10 * time.Second).C {
-			a.SaveMapping()
+		defer a.wg.Done()
+		ticker := time.NewTicker(10 * time.Second)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				a.SaveMapping()
+			case <-ctx.Done():
+				return
+			}
 		}
 	}()
 	return a
 }
 
+func (a *Anonymizer) Stop() {
+	a.cancel()
+	a.wg.Wait()
+}
+
 // SaveMapping writes the mapping from original to obfuscated strings to a file.
 // It is called by the anonymizer itself periodically, and should be called at
 // the end of the extraction run.

cmd/anonymizer/app/anonymizer/anonymizer_test.go

@@ -83,6 +83,7 @@ func TestNew(t *testing.T) {
 
 	file, err := os.CreateTemp(tempDir, "mapping.json")
 	require.NoError(t, err)
+	defer file.Close()
 
 	_, err = file.Write([]byte(`
 {
@@ -97,6 +98,7 @@ func TestNew(t *testing.T) {
 	require.NoError(t, err)
 
 	anonymizer := New(file.Name(), Options{}, nopLogger)
+	defer anonymizer.Stop()
 	assert.NotNil(t, anonymizer)
 }
 

cmd/anonymizer/app/writer/package_test.go

@@ -0,0 +1,14 @@
+// Copyright (c) 2024 The Jaeger Authors.
+// SPDX-License-Identifier: Apache-2.0
+
+package writer
+
+import (
+	"testing"
+
+	"github.com/jaegertracing/jaeger/pkg/testutils"
+)
+
+func TestMain(m *testing.M) {
+	testutils.VerifyGoLeaks(m)
+}

cmd/anonymizer/app/writer/writer.go

@@ -145,5 +145,6 @@ func (w *Writer) Close() {
 	w.capturedFile.Close()
 	w.anonymizedFile.WriteString("\n]\n")
 	w.anonymizedFile.Close()
+	w.anonymizer.Stop()
 	w.anonymizer.SaveMapping()
 }

cmd/anonymizer/app/writer/writer_test.go

@@ -53,8 +53,9 @@ func TestNew(t *testing.T) {
 			AnonymizedFile: tempDir + "/anonymized.json",
 			MappingFile:    tempDir + "/mapping.json",
 		}
-		_, err := New(config, nopLogger)
+		writer, err := New(config, nopLogger)
 		require.NoError(t, err)
+		defer writer.Close()
 	})
 
 	t.Run("CapturedFile does not exist", func(t *testing.T) {