From f4d3b0ca58ef731c3854f9a76134e5b45379ed04 Mon Sep 17 00:00:00 2001
From: teor <teor@riseup.net>
Date: Wed, 21 Jun 2023 12:30:14 +1000
Subject: [PATCH] Use Arc::into_inner() to avoid potential concurrency issues

---
 zebra-consensus/src/block.rs                   |  2 +-
 zebra-state/src/service.rs                     | 12 ++++--------
 zebra-state/src/service/non_finalized_state.rs |  2 ++
 3 files changed, 7 insertions(+), 9 deletions(-)

diff --git a/zebra-consensus/src/block.rs b/zebra-consensus/src/block.rs
index 3b694ac6773..970cf4118aa 100644
--- a/zebra-consensus/src/block.rs
+++ b/zebra-consensus/src/block.rs
@@ -280,7 +280,7 @@ where
             check::miner_fees_are_valid(&block, network, block_miner_fees)?;
 
             // Finally, submit the block for contextual verification.
-            let new_outputs = Arc::try_unwrap(known_utxos)
+            let new_outputs = Arc::into_inner(known_utxos)
                 .expect("all verification tasks using known_utxos are complete");
 
             let prepared_block = zs::SemanticallyVerifiedBlock {
diff --git a/zebra-state/src/service.rs b/zebra-state/src/service.rs
index 42e615516df..618b50ef1fe 100644
--- a/zebra-state/src/service.rs
+++ b/zebra-state/src/service.rs
@@ -270,7 +270,7 @@ impl Drop for ReadStateService {
 
         // TODO: move this into a try_shutdown() method
         if let Some(block_write_task) = self.block_write_task.take() {
-            if let Ok(block_write_task_handle) = Arc::try_unwrap(block_write_task) {
+            if let Some(block_write_task_handle) = Arc::into_inner(block_write_task) {
                 // We're the last database user, so we can tell it to shut down (blocking):
                 // - flushes the database to disk, and
                 // - drops the database, which cleans up any database tasks correctly.
@@ -1167,15 +1167,11 @@ impl Service<ReadRequest> for ReadStateService {
 
         if let Some(block_write_task) = block_write_task {
             if block_write_task.is_finished() {
-                match Arc::try_unwrap(block_write_task) {
+                if let Some(block_write_task) = Arc::into_inner(block_write_task) {
                     // We are the last state with a reference to this task, so we can propagate any panics
-                    Ok(block_write_task_handle) => {
-                        if let Err(thread_panic) = block_write_task_handle.join() {
-                            std::panic::resume_unwind(thread_panic);
-                        }
+                    if let Err(thread_panic) = block_write_task.join() {
+                        std::panic::resume_unwind(thread_panic);
                     }
-                    // We're not the last state, so we need to put it back
-                    Err(arc_block_write_task) => self.block_write_task = Some(arc_block_write_task),
                 }
             } else {
                 // It hasn't finished, so we need to put it back
diff --git a/zebra-state/src/service/non_finalized_state.rs b/zebra-state/src/service/non_finalized_state.rs
index bdcb1c10eb2..6cb9a2d447e 100644
--- a/zebra-state/src/service/non_finalized_state.rs
+++ b/zebra-state/src/service/non_finalized_state.rs
@@ -399,6 +399,8 @@ impl NonFinalizedState {
             // Pushing a block onto a Chain can launch additional parallel batches.
             // TODO: should we pass _scope into Chain::push()?
             scope.spawn_fifo(|_scope| {
+                // TODO: Replace with Arc::unwrap_or_clone() when it stabilises:
+                // https://github.com/rust-lang/rust/issues/93610
                 let new_chain = Arc::try_unwrap(new_chain)
                     .unwrap_or_else(|shared_chain| (*shared_chain).clone());
                 chain_push_result = Some(new_chain.push(contextual).map(Arc::new));