From 3b0ad9922f9f740a353a9a43ce602812fae7da89 Mon Sep 17 00:00:00 2001
From: mmsqe <mavis@crypto.com>
Date: Wed, 18 Dec 2024 09:46:11 +0800
Subject: [PATCH] fix: Better lock handling on close (backport #1024) (#1025)

Co-authored-by: Alexander Peters <alpe@users.noreply.github.com>
---
 nodedb.go      | 9 +++++----
 nodedb_test.go | 1 +
 2 files changed, 6 insertions(+), 4 deletions(-)

diff --git a/nodedb.go b/nodedb.go
index f1833b2eb..53db1908e 100644
--- a/nodedb.go
+++ b/nodedb.go
@@ -602,7 +602,7 @@ func (ndb *nodeDB) startPruning() {
 	for {
 		select {
 		case <-ndb.ctx.Done():
-			ndb.done <- struct{}{}
+			close(ndb.done)
 			return
 		default:
 			ndb.mtx.Lock()
@@ -1121,14 +1121,15 @@ func (ndb *nodeDB) traverseOrphans(prevVersion, curVersion int64, fn func(*Node)
 
 // Close the nodeDB.
 func (ndb *nodeDB) Close() error {
-	ndb.mtx.Lock()
-	defer ndb.mtx.Unlock()
-
 	ndb.cancel()
+
 	if ndb.opts.AsyncPruning {
 		<-ndb.done // wait for the pruning process to finish
 	}
 
+	ndb.mtx.Lock()
+	defer ndb.mtx.Unlock()
+
 	if ndb.batch != nil {
 		if err := ndb.batch.Close(); err != nil {
 			return err
diff --git a/nodedb_test.go b/nodedb_test.go
index 3e88c354f..b2d4796ab 100644
--- a/nodedb_test.go
+++ b/nodedb_test.go
@@ -444,4 +444,5 @@ func TestCloseNodeDB(t *testing.T) {
 	opts.AsyncPruning = true
 	ndb := newNodeDB(db, 0, opts, NewNopLogger())
 	require.NoError(t, ndb.Close())
+	require.NoError(t, ndb.Close()) // must not block or fail on second call
 }