Merge pull request #2782 from IntersectMBO/defragment_utxos

Add functionality to defragment address UTxOs

cardano_node_tests/defragment_utxos.py

@@ -0,0 +1,61 @@
+#!/usr/bin/env python3
+"""Defragment address UTxOs."""
+
+import argparse
+import logging
+import os
+import pathlib as pl
+import sys
+
+from cardano_clusterlib import clusterlib
+
+from cardano_node_tests.utils import defragment_utxos
+from cardano_node_tests.utils import helpers
+
+LOGGER = logging.getLogger(__name__)
+
+
+def get_args() -> argparse.Namespace:
+    """Get command line arguments."""
+    parser = argparse.ArgumentParser(description=__doc__.split("\n", maxsplit=1)[0])
+    parser.add_argument(
+        "-a",
+        "--address",
+        required=True,
+        help="Address",
+    )
+    parser.add_argument(
+        "-s",
+        "--skey-file",
+        type=helpers.check_file_arg,
+        help="Path to skey file",
+    )
+    return parser.parse_args()
+
+
+def main() -> int:
+    logging.basicConfig(
+        format="%(name)s:%(levelname)s:%(message)s",
+        level=logging.INFO,
+    )
+    args = get_args()
+
+    socket_env = os.environ.get("CARDANO_NODE_SOCKET_PATH")
+    if not socket_env:
+        LOGGER.error("The `CARDANO_NODE_SOCKET_PATH` environment variable is not set.")
+        return 1
+
+    state_dir = pl.Path(socket_env).parent
+    cluster_obj = clusterlib.ClusterLib(
+        state_dir=state_dir,
+        command_era=clusterlib.CommandEras.LATEST,
+    )
+    defragment_utxos.defragment(
+        cluster_obj=cluster_obj, address=args.address, skey_file=args.skey_file
+    )
+
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

cardano_node_tests/utils/defragment_utxos.py

@@ -0,0 +1,58 @@
+"""Defragment address UTxOs."""
+
+import logging
+import pathlib as pl
+
+from cardano_clusterlib import clusterlib
+
+LOGGER = logging.getLogger(__name__)
+
+
+def defragment(cluster_obj: clusterlib.ClusterLib, address: str, skey_file: pl.Path) -> None:
+    """Defragment address UTxOs."""
+    new_blocks = 3
+
+    loop = 1
+    utxos_len = -1
+    while True:
+        # Select UTxOs that are not locked and that contain only Lovelace
+        utxos_all = cluster_obj.g_query.get_utxo(address=address)
+        utxos_ids_excluded = {
+            f"{u.utxo_hash}#{u.utxo_ix}"
+            for u in utxos_all
+            if u.coin != clusterlib.DEFAULT_COIN or u.datum_hash
+        }
+        utxos = [u for u in utxos_all if f"{u.utxo_hash}#{u.utxo_ix}" not in utxos_ids_excluded]
+
+        prev_utxos_len, utxos_len = utxos_len, len(utxos)
+        if prev_utxos_len <= utxos_len and loop >= 2:
+            LOGGER.info("No more UTxOs to defragment.")
+            break
+        if utxos_len <= 10:
+            break
+
+        batch_size = min(100, utxos_len)
+        for b in range(1, utxos_len + 1, batch_size):
+            LOGGER.info(f"Defragmenting UTxOs: Running loop {loop}, batch {b}")
+            batch = utxos[b : b + batch_size]
+            tx_name = f"defrag_loop{loop}_batch{b}"
+
+            tx_output = cluster_obj.g_transaction.build_tx(
+                src_address=address,
+                tx_name=tx_name,
+                txins=batch,
+                change_address=address,
+            )
+            tx_signed_file = cluster_obj.g_transaction.sign_tx(
+                tx_body_file=tx_output.out_file,
+                tx_name=tx_name,
+                signing_key_files=[skey_file],
+            )
+            cluster_obj.g_transaction.submit_tx_bare(tx_file=tx_signed_file)
+
+        loop += 1
+
+        LOGGER.info(
+            f"Defragmenting UTxOs: Waiting for {new_blocks} new blocks before starting loop {loop}"
+        )
+        cluster_obj.wait_for_new_block(new_blocks=new_blocks)

cardano_node_tests/utils/testnet_cleanup.py

@@ -16,6 +16,7 @@
 from cardano_clusterlib import clusterlib
 
 from cardano_node_tests.utils import cluster_nodes
+from cardano_node_tests.utils import defragment_utxos
 
 LOGGER = logging.getLogger(__name__)
 
@@ -237,3 +238,8 @@ def _run(files: tp.List[pl.Path]) -> None:
     with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
         futures = [executor.submit(_run, f) for f in files_found]
         concurrent.futures.wait(futures)
+
+    # Defragment faucet address UTxOs
+    defragment_utxos.defragment(
+        cluster_obj=cluster_obj, address=faucet_payment.address, skey_file=faucet_payment.skey_file
+    )

pyproject.toml

@@ -65,6 +65,7 @@ sphinxemoji = "^0.3.1"
 
 [tool.poetry.scripts]
 testnet-cleanup = "cardano_node_tests.testnet_cleanup:main"
+defragment-utxos = "cardano_node_tests.defragment_utxos:main"
 prepare-cluster-scripts = "cardano_node_tests.prepare_cluster_scripts:main"
 split-topology = "cardano_node_tests.split_topology:main"
 cardano-cli-coverage = "cardano_node_tests.cardano_cli_coverage:main"