DL: remove data from DB on unsubscribe (#16786)

Co-authored-by: Kyle Altendorf <[email protected]>

Author: fchirica

chia/data_layer/data_layer.py

@@ -647,12 +647,12 @@ async def remove_subscriptions(self, store_id: bytes32, urls: List[str]) -> None
         async with self.subscription_lock:
             await self.data_store.remove_subscriptions(store_id, parsed_urls)
 
-    async def unsubscribe(self, tree_id: bytes32, retain_files: bool) -> None:
+    async def unsubscribe(self, tree_id: bytes32, retain_data: bool) -> None:
         subscriptions = await self.get_subscriptions()
         if tree_id not in (subscription.tree_id for subscription in subscriptions):
             raise RuntimeError("No subscription found for the given tree_id.")
         filenames: List[str] = []
-        if await self.data_store.tree_id_exists(tree_id) and not retain_files:
+        if await self.data_store.tree_id_exists(tree_id) and not retain_data:
             generation = await self.data_store.get_tree_generation(tree_id)
             all_roots = await self.data_store.get_roots_between(tree_id, 1, generation + 1)
             for root in all_roots:
@@ -663,6 +663,9 @@ async def unsubscribe(self, tree_id: bytes32, retain_files: bool) -> None:
         await self.wallet_rpc.dl_stop_tracking(tree_id)
         async with self.subscription_lock:
             await self.data_store.unsubscribe(tree_id)
+            if not retain_data:
+                await self.data_store.delete_store_data(tree_id)
+
         self.log.info(f"Unsubscribed to {tree_id}")
         for filename in filenames:
             file_path = self.server_files_location.joinpath(filename)

chia/data_layer/data_store.py

@@ -1073,13 +1073,35 @@ async def delete(
 
         return new_root
 
+    async def clean_node_table(self, writer: aiosqlite.Connection) -> None:
+        await writer.execute(
+            """
+            WITH RECURSIVE pending_nodes AS (
+                SELECT node_hash AS hash FROM root
+                WHERE status = ?
+                UNION ALL
+                SELECT n.left FROM node n
+                INNER JOIN pending_nodes pn ON n.hash = pn.hash
+                WHERE n.left IS NOT NULL
+                UNION ALL
+                SELECT n.right FROM node n
+                INNER JOIN pending_nodes pn ON n.hash = pn.hash
+                WHERE n.right IS NOT NULL
+            )
+            DELETE FROM node
+            WHERE hash NOT IN (SELECT hash FROM ancestors)
+            AND hash NOT IN (SELECT hash FROM pending_nodes)
+            """,
+            (Status.PENDING.value,),
+        )
+
     async def insert_batch(
         self,
         tree_id: bytes32,
         changelist: List[Dict[str, Any]],
         status: Status = Status.PENDING,
     ) -> Optional[bytes32]:
-        async with self.db_wrapper.writer():
+        async with self.db_wrapper.writer() as writer:
             old_root = await self.get_tree_root(tree_id)
             root_hash = old_root.node_hash
             if old_root.node_hash is None:
@@ -1146,6 +1168,8 @@ async def insert_batch(
                     "Didn't get the expected generation after batch update: "
                     f"Expected: {old_root.generation + 1}. Got: {new_root.generation}"
                 )
+
+            await self.clean_node_table(writer)
             return root.node_hash
 
     async def _get_one_ancestor(
@@ -1443,6 +1467,74 @@ async def remove_subscriptions(self, tree_id: bytes32, urls: List[str]) -> None:
                     },
                 )
 
+    async def delete_store_data(self, tree_id: bytes32) -> None:
+        async with self.db_wrapper.writer() as writer:
+            await self.clean_node_table(writer)
+            cursor = await writer.execute(
+                """
+                WITH RECURSIVE all_nodes AS (
+                    SELECT a.hash, n.left, n.right
+                    FROM ancestors AS a
+                    JOIN node AS n ON a.hash = n.hash
+                    WHERE a.tree_id = :tree_id
+                ),
+                pending_nodes AS (
+                    SELECT node_hash AS hash FROM root
+                    WHERE status = :status
+                    UNION ALL
+                    SELECT n.left FROM node n
+                    INNER JOIN pending_nodes pn ON n.hash = pn.hash
+                    WHERE n.left IS NOT NULL
+                    UNION ALL
+                    SELECT n.right FROM node n
+                    INNER JOIN pending_nodes pn ON n.hash = pn.hash
+                    WHERE n.right IS NOT NULL
+                )
+
+                SELECT hash, left, right
+                FROM all_nodes
+                WHERE hash NOT IN (SELECT hash FROM ancestors WHERE tree_id != :tree_id)
+                AND hash NOT IN (SELECT hash from pending_nodes)
+                """,
+                {"tree_id": tree_id, "status": Status.PENDING.value},
+            )
+            to_delete: Dict[bytes, Tuple[bytes, bytes]] = {}
+            ref_counts: Dict[bytes, int] = {}
+            async for row in cursor:
+                hash = row["hash"]
+                left = row["left"]
+                right = row["right"]
+                if hash in to_delete:
+                    prev_left, prev_right = to_delete[hash]
+                    assert prev_left == left
+                    assert prev_right == right
+                    continue
+                to_delete[hash] = (left, right)
+                if left is not None:
+                    ref_counts[left] = ref_counts.get(left, 0) + 1
+                if right is not None:
+                    ref_counts[right] = ref_counts.get(right, 0) + 1
+
+            await writer.execute("DELETE FROM ancestors WHERE tree_id == ?", (tree_id,))
+            await writer.execute("DELETE FROM root WHERE tree_id == ?", (tree_id,))
+            queue = [hash for hash in to_delete if ref_counts.get(hash, 0) == 0]
+            while queue:
+                hash = queue.pop(0)
+                if hash not in to_delete:
+                    continue
+                await writer.execute("DELETE FROM node WHERE hash == ?", (hash,))
+
+                left, right = to_delete[hash]
+                if left is not None:
+                    ref_counts[left] -= 1
+                    if ref_counts[left] == 0:
+                        queue.append(left)
+
+                if right is not None:
+                    ref_counts[right] -= 1
+                    if ref_counts[right] == 0:
+                        queue.append(right)
+
     async def unsubscribe(self, tree_id: bytes32) -> None:
         async with self.db_wrapper.writer() as writer:
             await writer.execute(

chia/rpc/data_layer_rpc_api.py

@@ -286,13 +286,13 @@ async def unsubscribe(self, request: Dict[str, Any]) -> EndpointResult:
         unsubscribe from singleton
         """
         store_id = request.get("id")
-        retain_files = request.get("retain", False)
+        retain_data = request.get("retain", False)
         if store_id is None:
             raise Exception("missing store id in request")
         if self.service is None:
             raise Exception("Data layer not created")
         store_id_bytes = bytes32.from_hexstr(store_id)
-        await self.service.unsubscribe(store_id_bytes, retain_files)
+        await self.service.unsubscribe(store_id_bytes, retain_data)
         return {}
 
     async def subscriptions(self, request: Dict[str, Any]) -> EndpointResult:

tests/core/data_layer/test_data_store.py

@@ -1428,3 +1428,185 @@ async def test_benchmark_batch_insert_speed(
             tree_id=tree_id,
             changelist=batch,
         )
+
+
+@pytest.mark.anyio
+async def test_delete_store_data(raw_data_store: DataStore) -> None:
+    tree_id = bytes32(b"\0" * 32)
+    tree_id_2 = bytes32(b"\0" * 31 + b"\1")
+    await raw_data_store.create_tree(tree_id=tree_id, status=Status.COMMITTED)
+    await raw_data_store.create_tree(tree_id=tree_id_2, status=Status.COMMITTED)
+    total_keys = 4
+    keys = [key.to_bytes(4, byteorder="big") for key in range(total_keys)]
+    batch1 = [
+        {"action": "insert", "key": keys[0], "value": keys[0]},
+        {"action": "insert", "key": keys[1], "value": keys[1]},
+    ]
+    batch2 = batch1.copy()
+    batch1.append({"action": "insert", "key": keys[2], "value": keys[2]})
+    batch2.append({"action": "insert", "key": keys[3], "value": keys[3]})
+    assert batch1 != batch2
+    await raw_data_store.insert_batch(tree_id, batch1, status=Status.COMMITTED)
+    await raw_data_store.insert_batch(tree_id_2, batch2, status=Status.COMMITTED)
+    keys_values_before = await raw_data_store.get_keys_values(tree_id_2)
+    async with raw_data_store.db_wrapper.reader() as reader:
+        result = await reader.execute("SELECT * FROM node")
+        nodes = await result.fetchall()
+        kv_nodes_before = {}
+        for node in nodes:
+            if node["key"] is not None:
+                kv_nodes_before[node["key"]] = node["value"]
+    assert [kv_nodes_before[key] for key in keys] == keys
+    await raw_data_store.delete_store_data(tree_id)
+    # Deleting from `node` table doesn't alter other stores.
+    keys_values_after = await raw_data_store.get_keys_values(tree_id_2)
+    assert keys_values_before == keys_values_after
+    async with raw_data_store.db_wrapper.reader() as reader:
+        result = await reader.execute("SELECT * FROM node")
+        nodes = await result.fetchall()
+        kv_nodes_after = {}
+        for node in nodes:
+            if node["key"] is not None:
+                kv_nodes_after[node["key"]] = node["value"]
+    for i in range(total_keys):
+        if i != 2:
+            assert kv_nodes_after[keys[i]] == keys[i]
+        else:
+            # `keys[2]` was only present in the first store.
+            assert keys[i] not in kv_nodes_after
+    assert not await raw_data_store.tree_id_exists(tree_id)
+    await raw_data_store.delete_store_data(tree_id_2)
+    async with raw_data_store.db_wrapper.reader() as reader:
+        async with reader.execute("SELECT COUNT(*) FROM node") as cursor:
+            row_count = await cursor.fetchone()
+            assert row_count is not None
+            assert row_count[0] == 0
+    assert not await raw_data_store.tree_id_exists(tree_id_2)
+
+
+@pytest.mark.anyio
+async def test_delete_store_data_multiple_stores(raw_data_store: DataStore) -> None:
+    # Make sure inserting and deleting the same data works
+    for repetition in range(2):
+        num_stores = 50
+        total_keys = 150
+        keys_deleted_per_store = 3
+        tree_ids = [bytes32(i.to_bytes(32, byteorder="big")) for i in range(num_stores)]
+        for tree_id in tree_ids:
+            await raw_data_store.create_tree(tree_id=tree_id, status=Status.COMMITTED)
+        original_keys = [key.to_bytes(4, byteorder="big") for key in range(total_keys)]
+        batches = []
+        for i in range(num_stores):
+            batch = [
+                {"action": "insert", "key": key, "value": key} for key in original_keys[i * keys_deleted_per_store :]
+            ]
+            batches.append(batch)
+
+        for tree_id, batch in zip(tree_ids, batches):
+            await raw_data_store.insert_batch(tree_id, batch, status=Status.COMMITTED)
+
+        for tree_index in range(num_stores):
+            async with raw_data_store.db_wrapper.reader() as reader:
+                result = await reader.execute("SELECT * FROM node")
+                nodes = await result.fetchall()
+
+            keys = {node["key"] for node in nodes if node["key"] is not None}
+            assert len(keys) == total_keys - tree_index * keys_deleted_per_store
+            keys_after_index = set(original_keys[tree_index * keys_deleted_per_store :])
+            keys_before_index = set(original_keys[: tree_index * keys_deleted_per_store])
+            assert keys_after_index.issubset(keys)
+            assert keys.isdisjoint(keys_before_index)
+            await raw_data_store.delete_store_data(tree_ids[tree_index])
+
+        async with raw_data_store.db_wrapper.reader() as reader:
+            async with reader.execute("SELECT COUNT(*) FROM node") as cursor:
+                row_count = await cursor.fetchone()
+                assert row_count is not None
+                assert row_count[0] == 0
+
+
+@pytest.mark.parametrize("common_keys_count", [1, 250, 499])
+@pytest.mark.anyio
+async def test_delete_store_data_with_common_values(raw_data_store: DataStore, common_keys_count: int) -> None:
+    tree_id_1 = bytes32(b"\x00" * 31 + b"\x01")
+    tree_id_2 = bytes32(b"\x00" * 31 + b"\x02")
+
+    await raw_data_store.create_tree(tree_id=tree_id_1, status=Status.COMMITTED)
+    await raw_data_store.create_tree(tree_id=tree_id_2, status=Status.COMMITTED)
+
+    key_offset = 1000
+    total_keys_per_store = 500
+    assert common_keys_count < key_offset
+    common_keys = {key.to_bytes(4, byteorder="big") for key in range(common_keys_count)}
+    unique_keys_1 = {
+        (key + key_offset).to_bytes(4, byteorder="big") for key in range(total_keys_per_store - common_keys_count)
+    }
+    unique_keys_2 = {
+        (key + (2 * key_offset)).to_bytes(4, byteorder="big") for key in range(total_keys_per_store - common_keys_count)
+    }
+
+    batch1 = [{"action": "insert", "key": key, "value": key} for key in common_keys.union(unique_keys_1)]
+    batch2 = [{"action": "insert", "key": key, "value": key} for key in common_keys.union(unique_keys_2)]
+
+    await raw_data_store.insert_batch(tree_id_1, batch1, status=Status.COMMITTED)
+    await raw_data_store.insert_batch(tree_id_2, batch2, status=Status.COMMITTED)
+
+    await raw_data_store.delete_store_data(tree_id_1)
+    async with raw_data_store.db_wrapper.reader() as reader:
+        result = await reader.execute("SELECT * FROM node")
+        nodes = await result.fetchall()
+
+    keys = {node["key"] for node in nodes if node["key"] is not None}
+    # Since one store got all its keys deleted, we're left only with the keys of the other store.
+    assert len(keys) == total_keys_per_store
+    assert keys.intersection(unique_keys_1) == set()
+    assert keys.symmetric_difference(common_keys.union(unique_keys_2)) == set()
+
+
+@pytest.mark.anyio
+async def test_delete_store_data_protects_pending_roots(raw_data_store: DataStore) -> None:
+    num_stores = 5
+    total_keys = 15
+    tree_ids = [bytes32(i.to_bytes(32, byteorder="big")) for i in range(num_stores)]
+    for tree_id in tree_ids:
+        await raw_data_store.create_tree(tree_id=tree_id, status=Status.COMMITTED)
+    original_keys = [key.to_bytes(4, byteorder="big") for key in range(total_keys)]
+    batches = []
+    keys_per_pending_root = 2
+
+    for i in range(num_stores - 1):
+        start_index = i * keys_per_pending_root
+        end_index = (i + 1) * keys_per_pending_root
+        batch = [{"action": "insert", "key": key, "value": key} for key in original_keys[start_index:end_index]]
+        batches.append(batch)
+    for tree_id, batch in zip(tree_ids, batches):
+        await raw_data_store.insert_batch(tree_id, batch, status=Status.PENDING)
+
+    tree_id = tree_ids[-1]
+    batch = [{"action": "insert", "key": key, "value": key} for key in original_keys]
+    await raw_data_store.insert_batch(tree_id, batch, status=Status.COMMITTED)
+
+    async with raw_data_store.db_wrapper.reader() as reader:
+        result = await reader.execute("SELECT * FROM node")
+        nodes = await result.fetchall()
+
+    keys = {node["key"] for node in nodes if node["key"] is not None}
+    assert keys == set(original_keys)
+
+    await raw_data_store.delete_store_data(tree_id)
+    async with raw_data_store.db_wrapper.reader() as reader:
+        result = await reader.execute("SELECT * FROM node")
+        nodes = await result.fetchall()
+
+    keys = {node["key"] for node in nodes if node["key"] is not None}
+    assert keys == set(original_keys[: (num_stores - 1) * keys_per_pending_root])
+
+    for index in range(num_stores - 1):
+        tree_id = tree_ids[index]
+        root = await raw_data_store.get_pending_root(tree_id)
+        assert root is not None
+        await raw_data_store.change_root_status(root, Status.COMMITTED)
+        kv = await raw_data_store.get_keys_values(tree_id=tree_id)
+        start_index = index * keys_per_pending_root
+        end_index = (index + 1) * keys_per_pending_root
+        assert {pair.key for pair in kv} == set(original_keys[start_index:end_index])