fix(discovery): move close mutex

ethpandaops · Oct 9, 2023 · 6aa8ac2 · 6aa8ac2
1 parent 3c0a14c
commit 6aa8ac2
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 10 deletions.
diff --git a/pkg/discovery/p2p/discovery/disc_v4.go b/pkg/discovery/p2p/discovery/disc_v4.go
@@ -303,9 +303,6 @@ func (d *DiscV4) OnNodeRecord(ctx context.Context, handler func(ctx context.Cont
 }
 
 func (l *ListenerV4) Close() error {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
 	if l.discovery != nil {
 		l.discovery.Close()
 	}
@@ -315,8 +312,13 @@ func (l *ListenerV4) Close() error {
 		l.localNode = nil
 	}
 
-	if l.conn != nil {
-		return l.conn.Close()
+	l.mu.Lock()
+	conn := l.conn
+	l.conn = nil
+	l.mu.Unlock()
+
+	if conn != nil {
+		return conn.Close()
 	}
 
 	return nil

diff --git a/pkg/discovery/p2p/discovery/disc_v5.go b/pkg/discovery/p2p/discovery/disc_v5.go
@@ -307,9 +307,6 @@ func (d *DiscV5) OnNodeRecord(ctx context.Context, handler func(ctx context.Cont
 }
 
 func (l *ListenerV5) Close() error {
-	l.mu.Lock()
-	defer l.mu.Unlock()
-
 	if l.discovery != nil {
 		l.discovery.Close()
 	}
@@ -319,8 +316,13 @@ func (l *ListenerV5) Close() error {
 		l.localNode = nil
 	}
 
-	if l.conn != nil {
-		return l.conn.Close()
+	l.mu.Lock()
+	conn := l.conn
+	l.conn = nil
+	l.mu.Unlock()
+
+	if conn != nil {
+		return conn.Close()
 	}
 
 	return nil