fix: Segfault/GIL state violation upon Python shutdown

.github/workflows/ci.yml

@@ -25,6 +25,8 @@ jobs:
     steps:
     - name: Checkout the revision
       uses: actions/checkout@v4
+      with:
+        submodules: recursive
     - name: Set up Rust toolchain
       uses: actions-rust-lang/setup-rust-toolchain@v1
       with:

.gitmodules

@@ -0,0 +1,3 @@
+[submodule "vendor/pyo3-async-runtimes"]
+	path = vendor/pyo3-async-runtimes
+	url = https://github.com/lablup/pyo3-async-runtimes.git

Cargo.toml

@@ -12,7 +12,8 @@ crate-type = ["cdylib"]
 [dependencies]
 etcd-client = "0.16.1"
 pyo3 = { version = "0.27.2", features = ["extension-module", "multiple-pymethods"] }
-pyo3-async-runtimes = { version = "0.27.0", features = ["attributes", "tokio-runtime"] }
+# Using patched version with explicit shutdown API
+pyo3-async-runtimes = { path = "vendor/pyo3-async-runtimes", features = ["attributes", "tokio-runtime"] }
 scopeguard = "1.2.0"
 tokio = { version = "1.46.1", features = ["sync"] }
 tokio-stream = "0.1.17"

etcd_client.pyi

@@ -358,3 +358,32 @@ class GRPCStatusCode(Enum):
 
     Unauthenticated = 16
     """The request does not have valid authentication credentials."""
+
+
+def cleanup_runtime() -> None:
+    """
+    Explicitly cleanup the tokio runtime.
+
+    This function signals the runtime to shutdown and waits for all tracked tasks
+    to complete. It should be called at the end of your main async function,
+    before the event loop shuts down.
+
+    Example:
+        ```python
+        from etcd_client import cleanup_runtime
+
+        async def main():
+            # Your etcd operations here
+            client = Client.connect(["localhost:2379"])
+            await client.put("key", "value")
+            # Cleanup before returning
+            cleanup_runtime()
+
+        asyncio.run(main())
+        ```
+
+    Note:
+        This is useful for ensuring clean shutdown and preventing GIL state
+        violations during Python interpreter finalization.
+    """
+    ...

pyproject.toml

@@ -43,6 +43,9 @@ dev = [
 requires = ["maturin>=1.7,<2.0"]
 build-backend = "maturin"
 
+[tool.maturin]
+python-source = "python"
+
 # === Tool Configuration ===
 
 [tool.ruff]

python/.gitignore

@@ -0,0 +1 @@
+*.so

python/etcd_client/__init__.py

@@ -0,0 +1,6 @@
+from .etcd_client import *  # noqa: F403
+from .etcd_client import cleanup_runtime  # noqa: F401
+
+__doc__ = etcd_client.__doc__  # noqa: F405
+if hasattr(etcd_client, "__all__"):  # noqa: F405
+    __all__ = etcd_client.__all__  # noqa: F405

scripts/run-etcd.sh

@@ -0,0 +1,23 @@
+ETCD_VER=v3.5.14
+
+rm -rf /tmp/etcd-data.tmp && mkdir -p /tmp/etcd-data.tmp && \
+  docker run -d \
+  -p 2379:2379 \
+  -p 2380:2380 \
+  --mount type=bind,source=/tmp/etcd-data.tmp,destination=/etcd-data \
+  --name etcd-gcr-${ETCD_VER} \
+  gcr.io/etcd-development/etcd:${ETCD_VER} \
+  /usr/local/bin/etcd \
+  --name s1 \
+  --data-dir /etcd-data \
+  --listen-client-urls http://0.0.0.0:2379 \
+  --advertise-client-urls http://0.0.0.0:2379 \
+  --listen-peer-urls http://0.0.0.0:2380 \
+  --initial-advertise-peer-urls http://0.0.0.0:2380 \
+  --initial-cluster s1=http://0.0.0.0:2380 \
+  --initial-cluster-token tkn \
+  --initial-cluster-state new \
+  --log-level info \
+  --logger zap \
+  --log-outputs stderr
+

src/client.rs

@@ -116,6 +116,7 @@ impl PyClient {
         result
     }
 
+    #[pyo3(signature = ())]
     fn __aenter__<'a>(&'a mut self, py: Python<'a>) -> PyResult<Bound<'a, PyAny>> {
         let endpoints = self.endpoints.clone();
         let connect_options = self.connect_options.clone();

src/communicator.rs

@@ -136,7 +136,6 @@ impl PyCommunicator {
         })
     }
 
-    // TODO: Implement and use the response types of `lease` type's methods
     fn lease_grant<'a>(&'a self, py: Python<'a>, ttl: i64) -> PyResult<Bound<'a, PyAny>> {
         let client = self.0.clone();
         future_into_py(py, async move {

src/condvar.rs

@@ -27,7 +27,7 @@ impl PyCondVar {
             while !*condition.lock().await {
                 inner.notified().await;
             }
-            Ok(())
+            Ok::<(), PyErr>(())
         })
     }
 
@@ -37,7 +37,7 @@ impl PyCondVar {
         future_into_py(py, async move {
             *condition.lock().await = true;
             inner.notify_waiters();
-            Ok(())
+            Ok::<(), PyErr>(())
         })
     }
 }

src/lib.rs

@@ -4,6 +4,7 @@ mod compare;
 mod condvar;
 mod error;
 mod lock_manager;
+mod runtime;
 mod txn;
 mod txn_response;
 mod watch;
@@ -70,6 +71,10 @@ mod etcd_client {
             py.get_type::<InvalidHeaderValueError>(),
         )?;
         m.add("EndpointError", py.get_type::<EndpointError>())?;
+
+        // Add runtime cleanup function
+        m.add_function(wrap_pyfunction!(crate::runtime::cleanup_runtime, m)?)?;
+
         Ok(())
     }
 }

src/runtime.rs

@@ -0,0 +1,25 @@
+use pyo3::prelude::*;
+
+/// Request graceful shutdown of the tokio runtime.
+///
+/// This should be called at the end of your main async function, before the event loop shuts down:
+///
+/// ```python
+/// import asyncio
+/// from etcd_client import cleanup_runtime
+///
+/// async def main():
+///     # Your etcd operations here
+///     ...
+///     # Cleanup before returning
+///     cleanup_runtime()
+///
+/// asyncio.run(main())
+/// ```
+///
+/// This function uses tokio's `shutdown_timeout()` to gracefully shut down all tasks,
+/// waiting up to 5 seconds for pending tasks to complete.
+#[pyfunction]
+pub fn cleanup_runtime() {
+    pyo3_async_runtimes::tokio::request_shutdown(5000);
+}

tests/conftest.py

@@ -27,47 +27,48 @@ def etcd_container():
     with DockerContainer(
         f"gcr.io/etcd-development/etcd:{ETCD_VER}",
         command=_etcd_command,
-    ).with_bind_ports("2379/tcp", 2379) as container:
+    ).with_exposed_ports(2379) as container:
         wait_for_logs(container, "ready to serve client requests")
-        yield
+        yield container
 
 
 @pytest.fixture
 async def etcd(etcd_container):
+    etcd_port = etcd_container.get_exposed_port(2379)
     etcd = AsyncEtcd(
-        addr=HostPortPair(host="127.0.0.1", port=2379),
+        addr=HostPortPair(host="127.0.0.1", port=etcd_port),
         namespace="test",
         scope_prefix_map={
             ConfigScopes.GLOBAL: "global",
             ConfigScopes.SGROUP: "sgroup/testing",
             ConfigScopes.NODE: "node/i-test",
         },
     )
-    try:
-        await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
-        await etcd.delete_prefix("", scope=ConfigScopes.SGROUP)
-        await etcd.delete_prefix("", scope=ConfigScopes.NODE)
-        yield etcd
-    finally:
-        await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
-        await etcd.delete_prefix("", scope=ConfigScopes.SGROUP)
-        await etcd.delete_prefix("", scope=ConfigScopes.NODE)
-        await etcd.close()
-        del etcd
+    async with etcd:
+        try:
+            await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
+            await etcd.delete_prefix("", scope=ConfigScopes.SGROUP)
+            await etcd.delete_prefix("", scope=ConfigScopes.NODE)
+            yield etcd
+        finally:
+            await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
+            await etcd.delete_prefix("", scope=ConfigScopes.SGROUP)
+            await etcd.delete_prefix("", scope=ConfigScopes.NODE)
 
 
 @pytest.fixture
 async def gateway_etcd(etcd_container):
+    etcd_port = etcd_container.get_exposed_port(2379)
     etcd = AsyncEtcd(
-        addr=HostPortPair(host="127.0.0.1", port=2379),
+        addr=HostPortPair(host="127.0.0.1", port=etcd_port),
         namespace="test",
         scope_prefix_map={
             ConfigScopes.GLOBAL: "",
         },
     )
-    try:
-        await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
-        yield etcd
-    finally:
-        await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
-        del etcd
+    async with etcd:
+        try:
+            await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)
+            yield etcd
+        finally:
+            await etcd.delete_prefix("", scope=ConfigScopes.GLOBAL)

tests/harness.py

@@ -9,6 +9,7 @@
 import functools
 import logging
 from collections import ChainMap, namedtuple
+from types import TracebackType
 from typing import (
     Any,
     AsyncGenerator,
@@ -165,8 +166,23 @@ def __init__(
             connect_options=self._connect_options,
         )
 
-    async def close(self):
-        pass  # for backward compatibility
+    async def open(self) -> None:
+        await self.etcd.__aenter__()
+
+    async def close(self) -> None:
+        await self.etcd.__aexit__(None, None, None)
+
+    async def __aenter__(self) -> "AsyncEtcd":
+        await self.etcd.__aenter__()
+        return self
+
+    async def __aexit__(
+        self,
+        exc_type: Optional[type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[TracebackType],
+    ) -> None:
+        await self.etcd.__aexit__(exc_type, exc_val, exc_tb)
 
     def _mangle_key(self, k: str) -> str:
         if k.startswith("/"):

tests/test.py → tests/test_function.py

@@ -246,6 +246,79 @@ async def _record_prefix():
     assert records_prefix[3].value == ""
 
 
+@pytest.mark.asyncio
+async def test_subprocess_segfault_reproduction(etcd_container) -> None:
+    """Test case to reproduce segfault when subprocess terminates quickly."""
+    import subprocess
+    import sys
+    import tempfile
+    import os
+    from pathlib import Path
+
+    # Create a script that will be run in subprocess
+    script_content = """
+import asyncio
+import sys
+
+from tests.harness import AsyncEtcd, ConfigScopes, HostPortPair
+
+async def main(etcd_port):
+    etcd = AsyncEtcd(
+        addr=HostPortPair(host="127.0.0.1", port=etcd_port),
+        namespace="test_subprocess",
+        scope_prefix_map={
+            ConfigScopes.GLOBAL: "global",
+        },
+    )
+
+    # Write a key and immediately exit
+    async with etcd:
+        await etcd.put("test_key", "test_value")
+
+if __name__ == "__main__":
+    etcd_port = int(sys.argv[1])
+    asyncio.run(main(etcd_port))
+"""
+    etcd_port = etcd_container.get_exposed_port(2379)
+
+    # Write the script to a temporary file
+    with tempfile.NamedTemporaryFile(mode="w", suffix=".py", delete=False) as f:
+        f.write(script_content)
+        script_path = f.name
+
+    # Get project root directory (parent of tests directory)
+    project_root = str(Path(__file__).parent.parent.resolve())
+
+    # Set up environment with PYTHONPATH
+    env = os.environ.copy()
+    env["PYTHONPATH"] = project_root
+
+    try:
+        # Run the subprocess 5 times to reproduce the segfault
+        for i in range(5):
+            result = subprocess.run(
+                [sys.executable, "-u", script_path, str(etcd_port)],
+                capture_output=True,
+                text=True,
+                timeout=10,
+                env=env,
+            )
+
+            # Check if the subprocess completed successfully
+            if result.returncode != 0:
+                print(f"Subprocess {i + 1} failed with return code {result.returncode}")
+                print(f"stderr: {result.stderr}")
+                print(f"stdout: {result.stdout}")
+
+            assert result.returncode == 0, (
+                f"Subprocess {i + 1} failed with return code {result.returncode}"
+            )
+
+    finally:
+        # Clean up the temporary script file
+        os.unlink(script_path)
+
+
 @pytest.mark.asyncio
 async def test_watch_once(etcd: AsyncEtcd) -> None:
     records = []