[object Object]

← back to Exo

wrote race-condition-free persistent NodeID-getting function

3ab56092899881ba81598e7e9f68cf17122616fb · 2025-07-23 20:18:56 +0100 · Andrei Cravtov

Files touched

Diff

commit 3ab56092899881ba81598e7e9f68cf17122616fb
Author: Andrei Cravtov <the.andrei.cravtov@gmail.com>
Date:   Wed Jul 23 20:18:56 2025 +0100

    wrote race-condition-free persistent NodeID-getting function
---
 .idea/externalDependencies.xml               |  2 +-
 justfile                                     |  2 +-
 rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi | 14 ++++-
 rust/exo_pyo3_bindings/src/pylibp2p/ident.rs | 32 ++++++++++-
 rust/exo_pyo3_bindings/tests/test_python.py  | 27 +++++++--
 shared/constants.py                          |  2 +
 shared/node_id.py                            | 51 +++++++++++++++++
 shared/pyproject.toml                        |  7 ++-
 shared/tests/test_node_id_persistence.py     | 85 ++++++++++++++++++++++++++++
 9 files changed, 212 insertions(+), 10 deletions(-)

diff --git a/.idea/externalDependencies.xml b/.idea/externalDependencies.xml
index c16deb13..60785b21 100644
--- a/.idea/externalDependencies.xml
+++ b/.idea/externalDependencies.xml
@@ -1,6 +1,6 @@
 <?xml version="1.0" encoding="UTF-8"?>
 <project version="4">
   <component name="ExternalDependencies">
-    <plugin id="com.insyncwithfoo.pyright" />
+    <plugin id="systems.fehn.intellijdirenv" />
   </component>
 </project>
\ No newline at end of file
diff --git a/justfile b/justfile
index 5865b22e..209cb5e5 100644
--- a/justfile
+++ b/justfile
@@ -17,7 +17,7 @@ lint-check:
     uv run ruff check master worker shared engines/*
 
 test:
-    uv run pytest master worker shared engines/* rust/exo_pyo3_bindings/tests
+    uv run pytest master worker shared engines/*
 
 check:
     uv run basedpyright --project pyproject.toml
diff --git a/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi b/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi
index 0cb78c74..f6e52b66 100644
--- a/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi
+++ b/rust/exo_pyo3_bindings/exo_pyo3_bindings.pyi
@@ -91,6 +91,10 @@ class Keypair:
         r"""
         TODO: documentation
         """
+    def to_peer_id(self) -> PeerId:
+        r"""
+        TODO: documentation
+        """
 
 class Multiaddr:
     r"""
@@ -143,6 +147,12 @@ class PeerId:
         r"""
         TODO: documentation
         """
-    def __repr__(self) -> builtins.str: ...
-    def __str__(self) -> builtins.str: ...
+    def __repr__(self) -> builtins.str:
+        r"""
+        TODO: documentation
+        """
+    def __str__(self) -> builtins.str:
+        r"""
+        TODO: documentation
+        """
 
diff --git a/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs b/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs
index 73239cca..39c01cf9 100644
--- a/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs
+++ b/rust/exo_pyo3_bindings/src/pylibp2p/ident.rs
@@ -3,7 +3,7 @@ use libp2p::identity::{ecdsa, Keypair};
 use libp2p::PeerId;
 use pyo3::prelude::{PyBytesMethods, PyModule, PyModuleMethods};
 use pyo3::types::PyBytes;
-use pyo3::{pyclass, pymethods, Bound, PyResult, Python};
+use pyo3::{pyclass, pymethods, Bound, PyObject, PyResult, Python};
 use pyo3_stub_gen::derive::{gen_stub_pyclass, gen_stub_pymethods};
 
 /// TODO: documentation...
@@ -76,6 +76,34 @@ impl PyKeypair {
         let bytes = self.0.to_protobuf_encoding().pyerr()?;
         Ok(PyBytes::new(py, &bytes))
     }
+
+    /// TODO: documentation
+    fn to_peer_id(&self) -> PyPeerId {
+        PyPeerId(self.0.public().to_peer_id())
+    }
+
+    // /// Hidden constructor for pickling support. TODO: figure out how to do pickling...
+    // #[gen_stub(skip)]
+    // #[new]
+    // fn py_new(bytes: Bound<'_, PyBytes>) -> PyResult<Self> {
+    //     Self::from_protobuf_encoding(bytes)
+    // }
+    //
+    // #[gen_stub(skip)]
+    // fn __setstate__(&mut self, state: Bound<'_, PyBytes>) -> PyResult<()> {
+    //     *self = Self::from_protobuf_encoding(state)?;
+    //     Ok(())
+    // }
+    //
+    // #[gen_stub(skip)]
+    // fn __getstate__<'py>(&self, py: Python<'py>) -> PyResult<Bound<'py, PyBytes>> {
+    //     self.to_protobuf_encoding(py)
+    // }
+    //
+    // #[gen_stub(skip)]
+    // pub fn __getnewargs__<'py>(&self, py: Python<'py>) -> PyResult<(Bound<'py, PyBytes>,)> {
+    //     Ok((self.to_protobuf_encoding(py)?,))
+    // }
 }
 
 /// TODO: documentation...
@@ -113,10 +141,12 @@ impl PyPeerId {
         self.0.to_base58()
     }
 
+    /// TODO: documentation
     fn __repr__(&self) -> String {
         format!("PeerId({})", self.to_base58())
     }
 
+    /// TODO: documentation
     fn __str__(&self) -> String {
         self.to_base58()
     }
diff --git a/rust/exo_pyo3_bindings/tests/test_python.py b/rust/exo_pyo3_bindings/tests/test_python.py
index d1408f45..1643c5a5 100644
--- a/rust/exo_pyo3_bindings/tests/test_python.py
+++ b/rust/exo_pyo3_bindings/tests/test_python.py
@@ -1,3 +1,7 @@
+import logging
+import multiprocessing
+import multiprocessing.queues
+import pickle
 import time
 from collections.abc import Awaitable
 from typing import Callable
@@ -48,7 +52,7 @@ async def test_discovery_callbacks() -> None:
     service.add_connected_callback(add_connected_callback)
     service.add_disconnected_callback(disconnected_callback)
 
-    for i in range(0, 10):
+    for i in range(0, 1):
         print(f"PYTHON: tick {i} of 10")
         time.sleep(1)
 
@@ -67,6 +71,21 @@ def disconnected_callback(e: ConnectionUpdate) -> None:
         f"PYTHON: Disconnected callback: {e.peer_id.__repr__()}, {e.connection_id.__repr__()}, {e.local_addr.__repr__()}, {e.send_back_addr.__repr__()}\n\n")
 
 
-async def foobar(a: Callable[[str], Awaitable[str]]):
-    abc = await a("")
-    pass
+# async def foobar(a: Callable[[str], Awaitable[str]]):
+#     abc = await a("")
+#     pass
+
+# def test_keypair_pickling() -> None:
+#     def subprocess_task(kp: Keypair, q: multiprocessing.queues.Queue[Keypair]):
+#         logging.info("a")
+#         assert q.get() == kp
+#         logging.info("b")
+#
+#
+#     kp = Keypair.generate_ed25519()
+#     q: multiprocessing.queues.Queue[Keypair] = multiprocessing.Queue()
+#
+#     p = multiprocessing.Process(target=subprocess_task, args=(kp, q))
+#     p.start()
+#     q.put(kp)
+#     p.join()
\ No newline at end of file
diff --git a/shared/constants.py b/shared/constants.py
index d187de03..61119538 100644
--- a/shared/constants.py
+++ b/shared/constants.py
@@ -9,6 +9,8 @@ EXO_WORKER_STATE = EXO_HOME / "worker_state.json"
 EXO_MASTER_LOG = EXO_HOME / "master.log"
 EXO_WORKER_LOG = EXO_HOME / "worker.log"
 
+EXO_NODE_ID_KEYPAIR = EXO_HOME / "node_id.keypair"
+
 EXO_WORKER_KEYRING_FILE = EXO_HOME / "worker_keyring"
 EXO_MASTER_KEYRING_FILE = EXO_HOME / "master_keyring"
 
diff --git a/shared/node_id.py b/shared/node_id.py
new file mode 100644
index 00000000..564a87a2
--- /dev/null
+++ b/shared/node_id.py
@@ -0,0 +1,51 @@
+import logging
+from multiprocessing import Lock
+from multiprocessing.synchronize import Lock as LockT
+from typing import Optional, TypedDict
+
+from exo_pyo3_bindings import Keypair
+
+from shared.constants import EXO_NODE_ID_KEYPAIR
+
+"""
+This file is responsible for concurrent race-free persistent node-ID retrieval.
+"""
+
+class _NodeIdGlobal(TypedDict):
+    file_lock: LockT
+    keypair: Optional[Keypair]
+
+_NODE_ID_GLOBAL: _NodeIdGlobal = {
+    "file_lock": Lock(),
+    "keypair": None,
+}
+
+def get_node_id_keypair() -> Keypair:
+    """
+    Obtains the :class:`Keypair` associated with this node-ID.
+    Obtain the :class:`PeerId` by from it.
+    """
+
+    # get from memory if we have it => read from file otherwise
+    if _NODE_ID_GLOBAL["keypair"] is not None:
+        return _NODE_ID_GLOBAL["keypair"]
+
+    # operate with cross-process lock to avoid race conditions
+    with _NODE_ID_GLOBAL["file_lock"]:
+        with open(EXO_NODE_ID_KEYPAIR, 'a+b') as f: # opens in append-mode => starts at EOF
+            # if non-zero EOF, then file exists => use to get node-ID
+            if f.tell() != 0:
+                f.seek(0) # go to start & read protobuf-encoded bytes
+                protobuf_encoded = f.read()
+
+                try: # if decoded successfully, save & return
+                    _NODE_ID_GLOBAL["keypair"] = Keypair.from_protobuf_encoding(protobuf_encoded)
+                    return _NODE_ID_GLOBAL["keypair"]
+                except RuntimeError as e: # on runtime error, assume corrupt file
+                    logging.warning(f"Encountered runtime error when trying to get keypair: {e}")
+
+        # if no valid credentials, create new ones and persist
+        with open(EXO_NODE_ID_KEYPAIR, 'w+b') as f:
+            _NODE_ID_GLOBAL["keypair"] = Keypair.generate_ed25519()
+            f.write(_NODE_ID_GLOBAL["keypair"].to_protobuf_encoding())
+            return _NODE_ID_GLOBAL["keypair"]
\ No newline at end of file
diff --git a/shared/pyproject.toml b/shared/pyproject.toml
index c4c5adeb..78920a59 100644
--- a/shared/pyproject.toml
+++ b/shared/pyproject.toml
@@ -15,7 +15,7 @@ dependencies = [
     "rustworkx>=0.16.0",
     "sqlmodel>=0.0.22",
     "sqlalchemy[asyncio]>=2.0.0",
-    "greenlet>=3.2.3"
+    "greenlet>=3.2.3",
 ]
 
 [build-system]
@@ -41,3 +41,8 @@ dev = [
     "pytest>=8.4.0",
     "pytest-asyncio>=1.0.0",
 ]
+
+[tool.pytest.ini_options]
+log_cli = true
+log_cli_level = "INFO"
+asyncio_mode = "auto"
diff --git a/shared/tests/test_node_id_persistence.py b/shared/tests/test_node_id_persistence.py
new file mode 100644
index 00000000..6f030b74
--- /dev/null
+++ b/shared/tests/test_node_id_persistence.py
@@ -0,0 +1,85 @@
+import contextlib
+import logging
+import os
+from multiprocessing import Event, Process, Queue, Semaphore
+from multiprocessing.queues import Queue as QueueT
+from multiprocessing.synchronize import Event as EventT
+from multiprocessing.synchronize import Semaphore as SemaphoreT
+from typing import Optional
+
+from pytest import LogCaptureFixture
+
+from shared.constants import EXO_NODE_ID_KEYPAIR
+from shared.node_id import get_node_id_keypair
+
+NUM_CONCURRENT_PROCS = 10
+
+def _get_keypair_concurrent(num_procs: int) -> bytes:
+    assert num_procs > 0
+
+    def subprocess_task(pid: int, sem: SemaphoreT, ev: EventT, queue: QueueT[bytes]) -> None:
+        # synchronise with parent process
+        logging.info(msg=f"SUBPROCESS {pid}: Started")
+        sem.release()
+
+        # wait to be told to begin simultaneous read
+        ev.wait()
+        logging.info(msg=f"SUBPROCESS {pid}: Reading start")
+        queue.put(get_node_id_keypair().to_protobuf_encoding())
+        logging.info(msg=f"SUBPROCESS {pid}: Reading end")
+
+        # notify master of finishing
+        sem.release()
+
+    sem = Semaphore(0)
+    ev = Event()
+    queue: QueueT[bytes] = Queue(maxsize=num_procs)
+
+    # make parent process wait for all subprocesses to start
+    logging.info(msg=f"PARENT: Starting {num_procs} subprocesses")
+    for i in range(num_procs):
+        Process(target=subprocess_task, args=(i + 1, sem, ev, queue)).start()
+    for _ in range(num_procs):
+        sem.acquire()
+
+    # start all the sub processes simultaneously
+    logging.info(msg="PARENT: Beginning read")
+    ev.set()
+
+    # wait until all subprocesses are done & read results
+    for _ in range(num_procs):
+        sem.acquire()
+
+    # check that the input/output order match, and that
+    # all subprocesses end up reading the same file
+    logging.info(msg="PARENT: Checking consistency")
+    keypair: Optional[bytes] = None
+    assert queue.qsize() > 0
+    while queue.qsize() > 0:
+        temp_keypair = queue.get()
+        if keypair is None:
+            keypair = temp_keypair
+        else:
+            assert keypair == temp_keypair
+    return keypair # pyright: ignore[reportReturnType]
+
+def _delete_if_exists(p: str | bytes | os.PathLike[str] | os.PathLike[bytes]):
+    with contextlib.suppress(OSError):
+        os.remove(p)
+
+def test_node_id_fetching(caplog: LogCaptureFixture):
+    reps = 10
+
+    # delete current file and write a new one
+    _delete_if_exists(EXO_NODE_ID_KEYPAIR)
+    kp = _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
+
+    with caplog.at_level(logging.CRITICAL): # supress logs
+        # make sure that continuous fetches return the same value
+        for _ in range(reps):
+            assert kp == _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
+
+        # make sure that after deleting, we are not fetching the same value
+        _delete_if_exists(EXO_NODE_ID_KEYPAIR)
+        for _ in range(reps):
+            assert kp != _get_keypair_concurrent(NUM_CONCURRENT_PROCS)
\ No newline at end of file

← 7a452c33 Fix tests  ·  back to Exo  ·  Add apply functions 56d35657 →