From 64c525b8d697028b02e2735c55a788e518dac6ec Mon Sep 17 00:00:00 2001
From: Matt Corallo <git@bluematt.me>
Date: Fri, 5 Feb 2021 13:20:14 -0300
Subject: [PATCH] Otherwise threads will try to continue modifying data
 structures while the dump process is being executed.

Flush CValidationInterface callbacks prior to destruction

Note that the CScheduler thread cant be running at this point,
it has already been stopped with the rest of the init threadgroup.
Thus, just calling any remaining loose callbacks during Shutdown()
is sane.
---
 src/init.cpp                | 13 +++++++++++++
 src/scheduler.cpp           |  9 +++++++++
 src/scheduler.h             |  3 +++
 src/test/test_pivx.cpp      |  1 +
 src/validationinterface.cpp |  4 ++++
 src/validationinterface.h   |  2 ++
 6 files changed, 32 insertions(+)

diff --git a/src/init.cpp b/src/init.cpp
index 7b19bcdea272c..66f858d86c93b 100644
--- a/src/init.cpp
+++ b/src/init.cpp
@@ -265,6 +265,19 @@ void PrepareShutdown()
         fFeeEstimatesInitialized = false;
     }
 
+    // FlushStateToDisk generates a SetBestChain callback, which we should avoid missing
+    FlushStateToDisk();
+
+    // After there are no more peers/RPC left to give us new data which may generate
+    // CValidationInterface callbacks, flush them...
+    GetMainSignals().FlushBackgroundCallbacks();
+
+    // Any future callbacks will be dropped. This should absolutely be safe - if
+    // missing a callback results in an unrecoverable situation, unclean shutdown
+    // would too. The only reason to do the above flushes is to let the wallet catch
+    // up with our current chain to avoid any strange pruning edge cases and make
+    // next startup faster by avoiding rescan.
+
     {
         LOCK(cs_main);
         if (pcoinsTip != NULL) {
diff --git a/src/scheduler.cpp b/src/scheduler.cpp
index 150c920283483..0a79f603b4f4e 100644
--- a/src/scheduler.cpp
+++ b/src/scheduler.cpp
@@ -171,3 +171,12 @@ void SingleThreadedSchedulerClient::AddToProcessQueue(std::function<void (void)>
     }
     MaybeScheduleProcessQueue();
 }
+
+void SingleThreadedSchedulerClient::EmptyQueue() {
+    bool should_continue = true;
+    while (should_continue) {
+        ProcessQueue();
+        LOCK(m_cs_callbacks_pending);
+        should_continue = !m_callbacks_pending.empty();
+    }
+}
diff --git a/src/scheduler.h b/src/scheduler.h
index 00a511bac7d2f..ada6370fce4f8 100644
--- a/src/scheduler.h
+++ b/src/scheduler.h
@@ -102,6 +102,9 @@ class SingleThreadedSchedulerClient {
 public:
     SingleThreadedSchedulerClient(CScheduler *pschedulerIn) : m_pscheduler(pschedulerIn) {}
     void AddToProcessQueue(std::function<void (void)> func);
+
+    // Processes all remaining queue members on the calling thread, blocking until queue is empty
+    void EmptyQueue();
 };
 
 #endif
diff --git a/src/test/test_pivx.cpp b/src/test/test_pivx.cpp
index 5ec8a5c0bbd9c..8537114b12854 100644
--- a/src/test/test_pivx.cpp
+++ b/src/test/test_pivx.cpp
@@ -86,6 +86,7 @@ TestingSetup::~TestingSetup()
         UnregisterNodeSignals(GetNodeSignals());
         threadGroup.interrupt_all();
         threadGroup.join_all();
+        GetMainSignals().FlushBackgroundCallbacks();
         GetMainSignals().UnregisterBackgroundSignalScheduler();
         UnloadBlockIndex();
         delete pcoinsTip;
diff --git a/src/validationinterface.cpp b/src/validationinterface.cpp
index 5c2741dbcf9f9..1af671456db4e 100644
--- a/src/validationinterface.cpp
+++ b/src/validationinterface.cpp
@@ -66,6 +66,10 @@ void CMainSignals::UnregisterBackgroundSignalScheduler() {
     m_internals.reset(nullptr);
 }
 
+void CMainSignals::FlushBackgroundCallbacks() {
+    m_internals->m_schedulerClient.EmptyQueue();
+}
+
 CMainSignals& GetMainSignals()
 {
     return g_signals;
diff --git a/src/validationinterface.h b/src/validationinterface.h
index 5ccdfc7e35c95..687c17222bf8f 100644
--- a/src/validationinterface.h
+++ b/src/validationinterface.h
@@ -60,6 +60,8 @@ class CMainSignals {
     void RegisterBackgroundSignalScheduler(CScheduler& scheduler);
     /** Unregister a CScheduler to give callbacks which should run in the background - these callbacks will now be dropped! */
     void UnregisterBackgroundSignalScheduler();
+    /** Call any remaining callbacks on the calling thread */
+    void FlushBackgroundCallbacks();
 
     void UpdatedBlockTip(const CBlockIndex *, const CBlockIndex *, bool fInitialDownload);
     void TransactionAddedToMempool(const CTransactionRef &ptxn);