From e06c8e83032d026bda45646ea14e59d47a9e8bf5 Mon Sep 17 00:00:00 2001
From: Jack Robison <jackrobison@lbry.io>
Date: Thu, 6 Jan 2022 12:36:33 -0500
Subject: [PATCH] update undo key to include the block hash

---
 lbry/wallet/server/block_processor.py |  4 ++--
 lbry/wallet/server/db/prefixes.py     | 24 ++++++++++++++----------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/lbry/wallet/server/block_processor.py b/lbry/wallet/server/block_processor.py
index 52965968c..827f654db 100644
--- a/lbry/wallet/server/block_processor.py
+++ b/lbry/wallet/server/block_processor.py
@@ -341,7 +341,7 @@ class BlockProcessor:
         def flush():
             self.db.write_db_state()
             if save_undo:
-                self.db.prefix_db.commit(self.height)
+                self.db.prefix_db.commit(self.height, self.tip)
             else:
                 self.db.prefix_db.unsafe_commit()
             self.clear_after_advance_or_reorg()
@@ -1559,7 +1559,7 @@ class BlockProcessor:
         self.db.last_flush_tx_count = self.db.fs_tx_count
 
         def rollback():
-            self.db.prefix_db.rollback(self.height + 1)
+            self.db.prefix_db.rollback(self.height + 1, reverted_block_hash)
             self.db.es_sync_height = self.height
             self.db.write_db_state()
             self.db.prefix_db.unsafe_commit()
diff --git a/lbry/wallet/server/db/prefixes.py b/lbry/wallet/server/db/prefixes.py
index ceb052ae8..f3baa71c4 100644
--- a/lbry/wallet/server/db/prefixes.py
+++ b/lbry/wallet/server/db/prefixes.py
@@ -1104,24 +1104,28 @@ class RepostedPrefixRow(PrefixRow):
         return cls.pack_key(reposted_claim_hash, tx_num, position), cls.pack_value(claim_hash)
 
 
+class UndoKey(NamedTuple):
+    height: int
+    block_hash: bytes
+
+
 class UndoPrefixRow(PrefixRow):
     prefix = DB_PREFIXES.undo.value
-    key_struct = struct.Struct(b'>Q')
+    key_struct = struct.Struct(b'>Q32s')
 
     key_part_lambdas = [
         lambda: b'',
-        struct.Struct(b'>Q').pack
+        struct.Struct(b'>Q').pack,
+        struct.Struct(b'>Q32s').pack
     ]
 
     @classmethod
-    def pack_key(cls, height: int):
-        return super().pack_key(height)
+    def pack_key(cls, height: int, block_hash: bytes):
+        return super().pack_key(height, block_hash)
 
     @classmethod
-    def unpack_key(cls, key: bytes) -> int:
-        assert key[:1] == cls.prefix
-        height, = cls.key_struct.unpack(key[1:])
-        return height
+    def unpack_key(cls, key: bytes) -> UndoKey:
+        return UndoKey(*super().unpack_key(key))
 
     @classmethod
     def pack_value(cls, undo_ops: bytes) -> bytes:
@@ -1132,8 +1136,8 @@ class UndoPrefixRow(PrefixRow):
         return data
 
     @classmethod
-    def pack_item(cls, height: int, undo_ops: bytes):
-        return cls.pack_key(height), cls.pack_value(undo_ops)
+    def pack_item(cls, height: int, block_hash: bytes, undo_ops: bytes):
+        return cls.pack_key(height, block_hash), cls.pack_value(undo_ops)
 
 
 class BlockHashPrefixRow(PrefixRow):