fix: Better lock handling on close (#1024)

cosmos · Dec 17, 2024 · 4f98c8a · 4f98c8a
1 parent b7bf354
commit 4f98c8a
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 4 deletions.
diff --git 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
@@ -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
 }