From 1e85686cf2df70ae7c2327f4f2787bba8c9649dc Mon Sep 17 00:00:00 2001
From: jholdstock <jholdstock@decred.org>
Date: Mon, 14 Oct 2024 09:11:33 +0100
Subject: [PATCH] jsonrpc: Stop background rescans (more) gracefully

Use the server waitgroup to ensure background rescans can return on
their own terms rather than being forcefully terminated early.

This is still not an ideal situation because in the case where the
server/wallet is being shut down, RescanFromHeight will return with an
"rpc connection closed" error (or similar) rather than returning because
its provided context was cancelled. However, that is a notable
improvement over the existing situation which could result in an
unhandled panic.
---
 internal/rpc/jsonrpc/methods.go | 33 ++++++++++++++++++++++++---------
 1 file changed, 24 insertions(+), 9 deletions(-)

diff --git a/internal/rpc/jsonrpc/methods.go b/internal/rpc/jsonrpc/methods.go
index 53b17f5c8..37afae274 100644
--- a/internal/rpc/jsonrpc/methods.go
+++ b/internal/rpc/jsonrpc/methods.go
@@ -2032,9 +2032,14 @@ func (s *Server) importPrivKey(ctx context.Context, icmd any) (any, error) {
 	}
 
 	if rescan {
-		// TODO: This is not synchronized with process shutdown and
-		// will cause panics when the DB is closed mid-transaction.
-		go w.RescanFromHeight(context.Background(), n, scanFrom)
+		// Rescan in the background rather than blocking the rpc request. Use
+		// the server waitgroup to ensure the rescan can return cleanly rather
+		// than being killed mid database transaction.
+		s.wg.Add(1)
+		go func() {
+			defer s.wg.Done()
+			_ = w.RescanFromHeight(context.Background(), n, scanFrom)
+		}()
 	}
 
 	return nil, nil
@@ -2084,9 +2089,14 @@ func (s *Server) importPubKey(ctx context.Context, icmd any) (any, error) {
 	}
 
 	if rescan {
-		// TODO: This is not synchronized with process shutdown and
-		// will cause panics when the DB is closed mid-transaction.
-		go w.RescanFromHeight(context.Background(), n, scanFrom)
+		// Rescan in the background rather than blocking the rpc request. Use
+		// the server waitgroup to ensure the rescan can return cleanly rather
+		// than being killed mid database transaction.
+		s.wg.Add(1)
+		go func() {
+			defer s.wg.Done()
+			_ = w.RescanFromHeight(context.Background(), n, scanFrom)
+		}()
 	}
 
 	return nil, nil
@@ -2130,9 +2140,14 @@ func (s *Server) importScript(ctx context.Context, icmd any) (any, error) {
 	}
 
 	if rescan {
-		// TODO: This is not synchronized with process shutdown and
-		// will cause panics when the DB is closed mid-transaction.
-		go w.RescanFromHeight(context.Background(), n, scanFrom)
+		// Rescan in the background rather than blocking the rpc request. Use
+		// the server waitgroup to ensure the rescan can return cleanly rather
+		// than being killed mid database transaction.
+		s.wg.Add(1)
+		go func() {
+			defer s.wg.Done()
+			_ = w.RescanFromHeight(context.Background(), n, scanFrom)
+		}()
 	}
 
 	return nil, nil