From 09011ebc7882c49d589f9af831541034dde88dd6 Mon Sep 17 00:00:00 2001
From: Marco Munizaga <git@marcopolo.io>
Date: Wed, 6 Jul 2022 10:01:46 -0700
Subject: [PATCH] Don't wait for a chan that will never close

---
 obs/stats_test.go | 16 ++++++++++++++++
 rcmgr_test.go     |  1 +
 trace.go          |  8 ++++----
 3 files changed, 21 insertions(+), 4 deletions(-)
 create mode 100644 obs/stats_test.go

diff --git a/obs/stats_test.go b/obs/stats_test.go
new file mode 100644
index 0000000..0236330
--- /dev/null
+++ b/obs/stats_test.go
@@ -0,0 +1,16 @@
+package obs_test
+
+import (
+	"testing"
+
+	rcmgr "github.com/libp2p/go-libp2p-resource-manager"
+	"github.com/libp2p/go-libp2p-resource-manager/obs"
+)
+
+func TestTraceReporterStartAndClose(t *testing.T) {
+	rcmgr, err := rcmgr.NewResourceManager(rcmgr.NewFixedLimiter(rcmgr.DefaultLimits.AutoScale()), rcmgr.WithTraceReporter(obs.StatsTraceReporter{}))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer rcmgr.Close()
+}
diff --git a/rcmgr_test.go b/rcmgr_test.go
index 6cf7e5f..de41cf6 100644
--- a/rcmgr_test.go
+++ b/rcmgr_test.go
@@ -1004,6 +1004,7 @@ func TestResourceManagerWithAllowlist(t *testing.T) {
 	if err != nil {
 		t.Fatal(err)
 	}
+	defer rcmgr.Close()
 
 	ableToGetAllowlist := GetAllowlist(rcmgr)
 	if ableToGetAllowlist == nil {
diff --git a/trace.go b/trace.go
index 396b8f9..ae871a0 100644
--- a/trace.go
+++ b/trace.go
@@ -19,7 +19,7 @@ type trace struct {
 
 	ctx    context.Context
 	cancel func()
-	closed chan struct{}
+	wg     sync.WaitGroup
 
 	mx            sync.Mutex
 	done          bool
@@ -232,7 +232,7 @@ func (t *trace) push(evt TraceEvt) {
 }
 
 func (t *trace) backgroundWriter(out io.WriteCloser) {
-	defer close(t.closed)
+	defer t.wg.Done()
 	defer out.Close()
 
 	gzOut := gzip.NewWriter(out)
@@ -315,7 +315,6 @@ func (t *trace) Start(limits Limiter) error {
 	}
 
 	t.ctx, t.cancel = context.WithCancel(context.Background())
-	t.closed = make(chan struct{})
 
 	if t.path != "" {
 		out, err := os.OpenFile(t.path, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644)
@@ -323,6 +322,7 @@ func (t *trace) Start(limits Limiter) error {
 			return nil
 		}
 
+		t.wg.Add(1)
 		go t.backgroundWriter(out)
 	}
 
@@ -350,7 +350,7 @@ func (t *trace) Close() error {
 	t.done = true
 	t.mx.Unlock()
 
-	<-t.closed
+	t.wg.Wait()
 	return nil
 }