Merge pull request #1 from egbertbouman/stats

Add message statistics and speedtest messages

Author: egbertbouman

.github/workflows/publish.yml

@@ -11,10 +11,22 @@ permissions:
 
 jobs:
   linux:
-    runs-on: ubuntu-latest
+    runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
-        target: [x86_64, x86, aarch64, armv7, s390x, ppc64le]
+        platform:
+          - runner: ubuntu-latest
+            target: x86_64
+          - runner: ubuntu-latest
+            target: x86
+          - runner: ubuntu-latest
+            target: aarch64
+          - runner: ubuntu-latest
+            target: armv7
+          - runner: ubuntu-latest
+            target: s390x
+          - runner: ubuntu-latest
+            target: ppc64le
     steps:
       - uses: actions/checkout@v3
       - uses: ./.github/actions/set-version
@@ -25,20 +37,24 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.target }}
-          args: --release --out dist -i 3.8 -i 3.9 -i 3.10 -i 3.11 -i 3.12
+          args: --release --out dist -i 3.9 -i 3.10 -i 3.11 -i 3.12 -i 3.13
           sccache: 'true'
           manylinux: auto
       - name: Upload wheels
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-linux-${{ matrix.platform.target }}
           path: dist
 
   windows:
-    runs-on: windows-latest
+    runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
-        target: [x64, x86]
+        platform:
+          - runner: windows-latest
+            target: x64
+          - runner: windows-latest
+            target: x86
     steps:
       - uses: actions/checkout@v3
       - uses: ./.github/actions/set-version
@@ -50,19 +66,23 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.target }}
-          args: --release --out dist -i 3.8 -i 3.9 -i 3.10 -i 3.11 -i 3.12
+          args: --release --out dist -i 3.9 -i 3.10 -i 3.11 -i 3.12 -i 3.13
           sccache: 'true'
       - name: Upload wheels
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-windows-${{ matrix.platform.target }}
           path: dist
 
   macos:
-    runs-on: macos-latest
+    runs-on: ${{ matrix.platform.runner }}
     strategy:
       matrix:
-        target: [x86_64, aarch64]
+        platform:
+          - runner: macos-13
+            target: x86_64
+          - runner: macos-14
+            target: aarch64
     steps:
       - uses: actions/checkout@v3
       - uses: ./.github/actions/set-version
@@ -73,12 +93,12 @@ jobs:
         uses: PyO3/maturin-action@v1
         with:
           target: ${{ matrix.target }}
-          args: --release --out dist -i 3.8 -i 3.9 -i 3.10 -i 3.11 -i 3.12
+          args: --release --out dist -i 3.9 -i 3.10 -i 3.11 -i 3.12 -i 3.13
           sccache: 'true'
       - name: Upload wheels
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-macos-${{ matrix.platform.target }}
           path: dist
 
   sdist:
@@ -92,9 +112,9 @@ jobs:
           command: sdist
           args: --out dist
       - name: Upload sdist
-        uses: actions/upload-artifact@v3
+        uses: actions/upload-artifact@v4
         with:
-          name: wheels
+          name: wheels-sdist
           path: dist
 
   release:
@@ -103,13 +123,11 @@ jobs:
     if: "startsWith(github.ref, 'refs/tags/')"
     needs: [linux, windows, macos, sdist]
     steps:
-      - uses: actions/download-artifact@v3
-        with:
-          name: wheels
+      - uses: actions/download-artifact@v4
       - name: Publish to PyPI
         uses: PyO3/maturin-action@v1
         env:
           MATURIN_PYPI_TOKEN: ${{ secrets.PYPI_TOKEN }}
         with:
           command: upload
-          args: --non-interactive --skip-existing *
+          args: --non-interactive --skip-existing wheels-*/*

.github/workflows/test.yml

@@ -16,10 +16,10 @@ jobs:
           toolchain: stable
           components: rustfmt
           target: x86_64-unknown-linux-gnu
-      - name: Setup Python 3.8
-        uses: actions/setup-python@v4
+      - name: Setup Python 3.9
+        uses: actions/setup-python@v5
         with:
-          python-version: 3.8
+          python-version: 3.9
       - name: Setup dependencies
         run: |
           pip install --upgrade pip

Cargo.toml

@@ -16,14 +16,14 @@ name = "rust_endpoint"
 crate-type = ["cdylib"]
 
 [dependencies]
-pyo3 = { version = "0.20.0", features = ["extension-module"] }
-tokio = { version = "1.34.0", features = ["full"] }
-env_logger = "0.10.1"
-log = "0.4.20"
-arc-swap = "1.6.0"
+pyo3 = { version = "0.23.4", features = ["extension-module"] }
+tokio = { version = "1.43.0", features = ["full"] }
+env_logger = "0.11.6"
+log = "0.4.25"
+arc-swap = "1.7.1"
 chacha20poly1305 = "0.10.1"
-socks5-proto = "0.4.0"
-socks5-server = "0.10.0"
-bytes = "1.5.0"
-rand = "0.8.5"
-map-macro = "0.2.6"
+socks5-proto = "0.4.1"
+socks5-server = "0.10.1"
+bytes = "1.10.0"
+rand = "0.9.0"
+map-macro = "0.3.0"

README.md

@@ -1,4 +1,4 @@
 # IPv8-rust-tunnels
-[![](https://img.shields.io/pypi/v/ipv8-rust-tunnels.svg?label=PyPI)](https://pypi.org/project/ipv8-rust-tunnels/)   ![Unit tests](https://github.com/egbertbouman/ipv8-rust-tunnels/actions/workflows/test.yml/badge.svg)
+[![](https://img.shields.io/pypi/v/ipv8-rust-tunnels.svg?label=PyPI)](https://pypi.org/project/ipv8-rust-tunnels/)   ![Unit tests](https://github.com/Tribler/ipv8-rust-tunnels/actions/workflows/test.yml/badge.svg)
 
 This module provides a set of performance enhancements to the `TunnelCommunity`, the anonymization layer used in [IPv8](https://github.com/Tribler/py-ipv8) and [Tribler](https://github.com/Tribler/tribler). It works by handling the tunnel data traffic in Rust, while letting the Python anonymization layer handle the tunnel control logic.

ipv8_rust_tunnels/endpoint.py

@@ -13,6 +13,8 @@
 import ipv8_rust_tunnels.rust_endpoint as rust
 
 from ipv8.messaging.anonymization.crypto import CryptoEndpoint
+from ipv8.messaging.interfaces.endpoint import EndpointListener
+from ipv8.messaging.interfaces.network_stats import NetworkStat
 from ipv8.messaging.interfaces.udp.endpoint import Endpoint, EndpointClosedException, UDPv4Address
 from ipv8.taskmanager import TaskManager
 from ipv8.util import succeed
@@ -65,6 +67,24 @@ def __init__(self, port: int = 0, ip: str = "0.0.0.0", worker_threads: int = 4)
 
         self.register_task('update_stats', self.update_stats, interval=1)
 
+    def add_prefix_listener(self, listener: EndpointListener, prefix: bytes) -> None:
+        """
+        Add an EndpointListener to our listeners, only triggers on packets with a specific prefix.
+
+        :raises: IllegalEndpointListenerError if the provided listener is not an EndpointListener
+        """
+        super().add_prefix_listener(listener, prefix)
+        if self.rust_ep.is_open():
+            self.rust_ep.set_prefixes(list(self._prefix_map.keys()))
+
+    def remove_listener(self, listener: EndpointListener) -> None:
+        """
+        Remove a listener from our listeners, if it is registered.
+        """
+        super().remove_listener(listener)
+        if self.rust_ep.is_open():
+            self.rust_ep.set_prefixes(list(self._prefix_map.keys()))
+
     def update_stats(self) -> None:
         """
         Updates the statistics of the routing objects using the most recent data from Rust.
@@ -78,6 +98,25 @@ def update_stats(self) -> None:
         for exit_socket in self.exit_sockets.values():
             self.rust_ep.update_exit_stats(exit_socket.circuit_id, exit_socket)
 
+    def get_statistics(self, prefix: bytes) -> dict[int, NetworkStat]:
+        """
+        Get the message statistics per message identifier for the given prefix.
+        """
+        result = {}
+        for msg_id, counters in self.rust_ep.get_message_statistics(prefix).items():
+            stat = result[msg_id] = NetworkStat(msg_id)
+            stat.num_up = counters[0]
+            stat.num_down = counters[2]
+            stat.bytes_up = counters[1]
+            stat.bytes_down = counters[3]
+        return result
+
+    def enable_community_statistics(self, community_prefix: bytes, enabled: bool) -> None:
+        """
+        Start tracking stats for packets with the given prefix.
+        """
+        pass
+
     def setup_tunnels(self, tunnel_community: TunnelCommunity, settings: TunnelSettings) -> None:
         """
         Set up the TunnelCommunity.
@@ -93,6 +132,7 @@ def apply_settings(self) -> None:
         """
         if self.prefix and self.settings and self.is_open():
             self.rust_ep.set_prefix(self.prefix)
+            self.rust_ep.set_prefixes(list(self._prefix_map.keys()))
             self.rust_ep.set_max_relay_early(self.settings.max_relay_early)
             self.rust_ep.set_peer_flags(self.settings.peer_flags)
 
@@ -210,11 +250,12 @@ def reset_byte_counters(self) -> None:
         self.bytes_up = 0
         self.bytes_down = 0
 
-    def run_speedtest(self, target_addr: str, associate_port: int, num_packets: int, request_size: int,
-                       response_size: int, timeout_ms: int, window_size: int, callback: Callable) -> None:
+    def run_speedtest(self, circuit_id: int, test_time: int, request_size: int,
+                      response_size: int, target_rtt: int, callback: Callable, callback_interval: int = 0) -> None:
         """
-        Perform a TunnelCommunity speedtest. Connects to an existing UDP associate
-        port and sends test messages to a given target address.
+        Perform a TunnelCommunity speedtest.
         """
-        self.rust_ep.run_speedtest(target_addr, associate_port, num_packets, request_size,
-                                   response_size, timeout_ms, window_size, callback)
+        def callback_threadsafe(*args):
+            self.loop.call_soon_threadsafe(callback, *args)
+        return self.rust_ep.run_speedtest(circuit_id, test_time, request_size,
+                                          response_size, target_rtt, callback_threadsafe, callback_interval)

ipv8_rust_tunnels/tests/test_tunnel_community.py

@@ -1,7 +1,12 @@
+import asyncio
 import unittest
+from unittest.mock import Mock
 
 from ipv8_rust_tunnels.endpoint import RustEndpoint
 
+from ipv8.messaging.anonymization.community import CIRCUIT_TYPE_RP_DOWNLOADER
+from ipv8.messaging.anonymization.tunnel import PEER_FLAG_SPEED_TEST
+from ipv8.messaging.interfaces.endpoint import EndpointListener
 from ipv8.messaging.interfaces.udp.endpoint import UDPv4Address
 from ipv8.test.messaging.anonymization.test_community import TestTunnelCommunity
 from ipv8.test.messaging.anonymization.test_hiddenservices import TestHiddenServices
@@ -26,6 +31,12 @@ def create_node(org, *args, **kwargs):  # noqa: ANN201, ANN002, ANN001, D103
     ep.remove_listener(ipv8.overlay)
     ep.add_prefix_listener(ipv8.overlay, overlay.get_prefix())
 
+    # Some unittests use prefixes that we don't want blocked. So, we add a fake EndpointListener.
+    mock = Mock()
+    mock.__class__ = EndpointListener
+    ep.add_prefix_listener(mock, b'\x00' * 22)
+    ep.add_prefix_listener(mock, b'\x00\x01' + b'\x00' * 20)
+
     overlay.circuits = ep.circuits
     overlay.relay_from_to = ep.relays
     overlay.exit_sockets = ep.exit_sockets
@@ -41,17 +52,70 @@ def replace(old_func, new_func):  # noqa: ANN001, D103, ANN201
     return lambda *args, org=old_func, **kwargs: new_func(org, *args, **kwargs)
 
 
+async def new_test_test_request(self) -> None:
+    """
+    Check if sending test-request messages works as expected.
+    """
+    self.add_node_to_experiment(self.create_node())
+    self.settings(1).peer_flags |= {PEER_FLAG_SPEED_TEST}
+    await self.introduce_nodes()
+    circuit = self.overlay(0).create_circuit(2, exit_flags=[PEER_FLAG_SPEED_TEST])
+    await circuit.ready
+
+    callback = Mock()
+    self.overlay(0).endpoint.run_speedtest(circuit.circuit_id, 10, 10, 10, 1, callback)
+    await asyncio.sleep(.1)
+
+    callback.assert_called_once()
+    self.assertEqual(len(callback.call_args_list), 1)
+    msg_stats = list(callback.call_args_list[0].args[0].values())
+    self.assertTrue(sum([stat[1] for stat in msg_stats]) > 0)
+    self.assertTrue(sum([stat[3] for stat in msg_stats]) > 0)
+
+
+async def new_test_test_request_e2e(self, *args) -> None:
+    """
+    Check if sending test-request messages over an e2e circuit works as expected.
+    """
+    future = asyncio.Future()
+
+    self.overlay(0).join_swarm(self.service, 1, future.set_result, seeding=False)
+    self.overlay(2).join_swarm(self.service, 1, future.set_result)
+    self.overlay(2).settings.peer_flags.add(PEER_FLAG_SPEED_TEST)
+
+    await self.introduce_nodes()
+    await self.create_intro(2, self.service)
+    await self.assign_exit_node(0)
+
+    await self.overlay(0).do_peer_discovery()
+    await self.deliver_messages()
+
+    await future
+
+    circuit, = self.overlay(0).find_circuits(ctype=CIRCUIT_TYPE_RP_DOWNLOADER)
+    callback = Mock()
+    self.overlay(0).endpoint.run_speedtest(circuit.circuit_id, 10, 3, 6, 1, callback)
+    await asyncio.sleep(.1)
+
+    callback.assert_called_once()
+    self.assertEqual(len(callback.call_args_list), 1)
+    msg_stats = list(callback.call_args_list[0].args[0].values())
+    self.assertTrue(sum([stat[1] for stat in msg_stats]) > 0)
+    self.assertTrue(sum([stat[3] for stat in msg_stats]) > 0)
+
+
 TestTunnelCommunity.create_node = replace(TestTunnelCommunity.create_node, create_node)
 TestTunnelCommunity.setUp = replace(TestTunnelCommunity.setUp, set_up)
+TestTunnelCommunity.test_test_request = new_test_test_request
 TestTunnelCommunity.test_tunnel_unicode_destination = \
     unittest.skip("not available in RustEndpoint")(TestTunnelCommunity.test_tunnel_unicode_destination)
 
 TestHiddenServices.create_node = replace(TestHiddenServices.create_node, create_node)
 TestHiddenServices.setUp = replace(TestHiddenServices.setUp, set_up)
+TestHiddenServices.test_test_request_e2e = new_test_test_request_e2e
 TestHiddenServices.test_dht_lookup_with_counterparty = \
     unittest.skip("not available in RustEndpoint")(TestHiddenServices.test_dht_lookup_with_counterparty)
 
-
 if __name__ == '__main__':
     runner = unittest.TextTestRunner()
     runner.run(unittest.makeSuite(TestTunnelCommunity))

pyproject.toml

@@ -5,7 +5,7 @@ build-backend = "maturin"
 [project]
 name = "ipv8-rust-tunnels"
 description = "IPv8 tunnel performance enhancements"
-requires-python = ">=3.8"
+requires-python = ">=3.9"
 classifiers = [
     "Programming Language :: Rust",
     "Programming Language :: Python :: Implementation :: CPython",