feat: Automate tokio runtime cleanup via reference counting (#17)

Author: achimnol

.github/workflows/ci.yml

@@ -33,6 +33,9 @@ jobs:
       with:
         toolchain: stable
         override: true
+        # Include Python version in cache key to prevent cross-contamination
+        # between GIL-enabled (3.14) and free-threaded (3.14t) builds
+        cache-key: py${{ matrix.python-version }}
     - name: Install protobuf compiler
       uses: arduino/setup-protoc@v3
       with:
@@ -43,6 +46,11 @@ jobs:
       with:
         enable-cache: true
         python-version: ${{ matrix.python-version }}
+    - name: Check python build
+      run: |
+        uv run --no-project python -c 'import sys; print(sys.prefix); print(sys.version_info)'
+        uv run --no-project python -c 'import sysconfig; flag=sysconfig.get_config_var("Py_GIL_DISABLED"); print(f"Py_GIL_DISABLED={flag}")'
+        uv run --no-project python -c 'import sys; print(f"{sys.abiflags=}")'
     - name: Install dependencies and build the package
       run: |
         uv sync --locked --all-extras --no-install-project

README.md

@@ -15,7 +15,7 @@ pip install etcd_client
 
 ```python
 from etcd_client import EtcdClient
-etcd = EtcdClient(['http:://127.0.0.1:2379'])
+etcd = EtcdClient(['http://127.0.0.1:2379'])
 ```
 
 Actual connection establishment with Etcd's gRPC channel will be done when you call `EtcdClient.connect()`.
@@ -28,28 +28,9 @@ async def main():
         print(bytes(value).decode())  # testvalue
 ```
 
-### Cleanup on shutdown
+### Working with key prefixes
 
-To prevent segfaults or GIL state violations during Python interpreter shutdown, you should call `cleanup_runtime()` at the end of your main async function before the event loop shuts down:
-
-```python
-from etcd_client import EtcdClient, cleanup_runtime
-
-async def main():
-    etcd = EtcdClient(['http://127.0.0.1:2379'])
-    async with etcd.connect() as communicator:
-        await communicator.put('testkey'.encode(), 'testvalue'.encode())
-        value = await communicator.get('testkey'.encode())
-        print(bytes(value).decode())
-    # Cleanup the tokio runtime before returning
-    cleanup_runtime()
-
-asyncio.run(main())
-```
-
-This function signals the internal tokio runtime to shut down gracefully, waiting up to 5 seconds for pending tasks to complete.
-
-`EtcdCommunicator.get_prefix(prefix)` will return a tuple of list containing all key-values with given key prefix.
+`EtcdCommunicator.get_prefix(prefix)` returns a list of key-value pairs matching the given prefix.
 
 ```python
 async def main():
@@ -69,9 +50,39 @@ async def main():
             print([bytes(v).decode() for v in resp])
 ```
 
+## Automatic runtime cleanup
+
+The tokio runtime is automatically cleaned up when the last client context exits. In most cases, no explicit cleanup is needed:
+
+```python
+import asyncio
+from etcd_client import EtcdClient
+
+async def main():
+    etcd = EtcdClient(['http://127.0.0.1:2379'])
+    async with etcd.connect() as communicator:
+        await communicator.put('testkey'.encode(), 'testvalue'.encode())
+        value = await communicator.get('testkey'.encode())
+        print(bytes(value).decode())
+    # Runtime automatically cleaned up when context exits
+
+asyncio.run(main())
+```
+
+The library uses reference counting to track active client contexts. When the last context exits, the tokio runtime is gracefully shut down, waiting up to 5 seconds for pending tasks to complete. If you create new clients after this, the runtime is automatically re-initialized.
+
+For advanced use cases requiring explicit control, `cleanup_runtime()` is available:
+
+```python
+from etcd_client import cleanup_runtime
+
+# Force cleanup at a specific point (usually not needed)
+cleanup_runtime()
+```
+
 ## Operating with Etcd lock
 
-Just like `EtcdClient.connect()`, you can easilly use etcd lock by calling `EtcdClient.with_lock(lock_opts)`.
+Just like `EtcdClient.connect()`, you can easily use etcd lock by calling `EtcdClient.with_lock(lock_opts)`.
 
 ```python
 async def first():
@@ -98,18 +109,11 @@ async with etcd.connect() as communicator:
 await asyncio.gather(first(), second())  # first: testvalue | second: testvalue
 ```
 
-Adding `timeout` parameter to `EtcdClient.with_lock()` call will add a timeout to lock acquiring process.
+### Lock timeout
 
-```python
-async def first():
-    async with etcd.with_lock(
-        EtcdLockOption(
-            lock_name='foolock'.encode(),
-        )
-    ) as communicator:
-        value = await communicator.get('testkey'.encode())
-        print('first:', bytes(value).decode(), end=' | ')
+Adding `timeout` parameter to `EtcdLockOption` will add a timeout to the lock acquiring process.
 
+```python
 async def second():
     await asyncio.sleep(0.1)
     async with etcd.with_lock(
@@ -120,13 +124,11 @@ async def second():
     ) as communicator:
         value = await communicator.get('testkey'.encode())
         print('second:', bytes(value).decode())
-
-async with etcd.connect() as communicator:
-    await communicator.put('testkey'.encode(), 'testvalue'.encode())
-await asyncio.gather(first(), second())  # first: testvalue | second: testvalue
 ```
 
-Adding `ttl` parameter to `EtcdClient.with_lock()` call will force lock to be released after given seconds.
+### Lock TTL
+
+Adding `ttl` parameter to `EtcdLockOption` will force the lock to be released after the given seconds.
 
 ```python
 async def first():
@@ -154,7 +156,7 @@ for task in done:
 
 ## Watch
 
-You can watch changes on key with `EtcdCommunicator.watch(key)`.
+You can watch changes on a key with `EtcdCommunicator.watch(key)`.
 
 ```python
 async def watch():
@@ -179,7 +181,9 @@ await asyncio.gather(watch(), update())
 # WatchEventType.PUT 5
 ```
 
-Watching changes on keys with specific prefix can be also done by `EtcdCommunicator.watch_prefix(key_prefix)`.
+### Watch with prefix
+
+Watching changes on keys with a specific prefix can be done with `EtcdCommunicator.watch_prefix(key_prefix)`.
 
 ```python
 async def watch():
@@ -204,11 +208,11 @@ await asyncio.gather(watch(), update())
 
 ## Transaction
 
-You can run etcd transaction by calling `EtcdCommunicator.txn(txn)`.
+You can run etcd transactions by calling `EtcdCommunicator.txn(txn)`.
 
 ### Constructing compares
 
-Constructing compare operations can be done by comparing `Compare` instance.
+Constructing compare operations can be done using the `Compare` class.
 
 ```python
 from etcd_client import Compare, CompareOp
@@ -218,7 +222,7 @@ compares = [
 ]
 ```
 
-### Executing transaction calls
+### Executing transactions
 
 ```python
 async with etcd.connect() as communicator:
@@ -232,29 +236,26 @@ async with etcd.connect() as communicator:
     ]
 
     res = await communicator.txn(Txn().when(compares).and_then([TxnOp.get('successkey'.encode())]))
-    print(res) # TODO: Need to write response type bindings.
+    print(res)  # TODO: Need to write response type bindings.
 ```
 
 ## How to build
 
-### Prerequisite
+### Prerequisites
 
-* The Rust development environment (the 2021 edition or later) using [`rustup`](https://rustup.rs/) or your package manager
+* The Rust development environment (2021 edition or later) using [`rustup`](https://rustup.rs/) or your package manager
 * The Python development environment (3.10 or later) using [`pyenv`](https://github.com/pyenv/pyenv#installation) or your package manager
 
-### Build instruction
+### Build instructions
 
-First, create a virtualenv (either using the standard venv package, pyenv, or
-whatever your favorite).  Then, install the PEP-517 build toolchain and run it.
+First, create a virtualenv (using the standard venv package, pyenv, or your preferred tool). Then, install the PEP-517 build toolchain and run it.
 
 ```shell
 pip install -U pip build setuptools
 python -m build --sdist --wheel
 ```
 
-It will automatically install build dependencies like
-[`maturin`](https://github.com/PyO3/maturin) and build the wheel and source
-distributions under the `dist/` directory.
+This will automatically install build dependencies like [`maturin`](https://github.com/PyO3/maturin) and build the wheel and source distributions under the `dist/` directory.
 
 ## How to develop and test
 
@@ -279,42 +280,18 @@ uv run maturin develop  # Builds and installs the Rust extension
 This project uses ruff for linting/formatting and mypy for type checking:
 
 ```bash
-# Format Python code
-make fmt-py
-
-# Lint Python code
-make lint-py
-
-# Auto-fix Python issues (format + fixable lints)
-make fix-py
-
-# Type check Python code
-make typecheck
-
-# Auto-fix Rust issues (format + fixable clippy lints)
-make fix-rust
-
-# Auto-fix all issues (Python + Rust)
-make fix
-
-# Format all code (Python + Rust)
-make fmt
-
-# Lint all code (Python + Rust)
-make lint
-
-# Run all checks (Python + Rust)
-make check
+make fmt       # Format all code (Python + Rust)
+make lint      # Lint all code (Python + Rust)
+make fix       # Auto-fix all issues (Python + Rust)
+make typecheck # Type check Python code
+make check     # Run all checks
 ```
 
 ### Running tests
 
 ```bash
-# Run tests using uv
-make test
-
-# Or directly with uv
-uv run pytest
+make test      # Run tests using uv
+uv run pytest  # Or directly with uv
 
-# The tests use testcontainers to automatically spin up etcd
+# Tests use testcontainers to automatically spin up etcd
 ```

etcd_client.pyi

@@ -360,30 +360,62 @@ class GRPCStatusCode(Enum):
     """The request does not have valid authentication credentials."""
 
 
+def active_context_count() -> int:
+    """
+    Get the number of currently active client contexts.
+
+    Returns the count of client context managers currently in use (inside
+    `async with` blocks). This is useful for debugging and testing the
+    automatic cleanup behavior.
+
+    Returns:
+        The number of active client contexts. Returns 0 when no clients
+        are in an active context manager.
+
+    Example:
+        ```python
+        from etcd_client import Client, active_context_count
+
+        client = Client(["localhost:2379"])
+        print(active_context_count())  # 0
+
+        async with client.connect():
+            print(active_context_count())  # 1
+
+        print(active_context_count())  # 0
+        ```
+    """
+    ...
+
+
 def cleanup_runtime() -> None:
     """
     Explicitly cleanup the tokio runtime.
 
+    In most cases, the runtime is automatically cleaned up when the last
+    client context exits. This function is provided for cases where explicit
+    control is needed, such as when using the client without a context manager.
+
     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.
+    to complete (up to 5 seconds). After shutdown, the runtime will be lazily
+    re-initialized if new client operations are performed.
 
     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()
+            async with client.connect():
+                await client.put("key", "value")
+            # Runtime is automatically cleaned up when context exits
+            # Explicit call is usually not needed
 
         asyncio.run(main())
         ```
 
     Note:
-        This is useful for ensuring clean shutdown and preventing GIL state
-        violations during Python interpreter finalization.
+        This function is idempotent - calling it multiple times or when the
+        runtime is already shut down is safe and has no effect.
     """
     ...

pyproject.toml

@@ -22,11 +22,11 @@ classifiers = [
   "Programming Language :: Python :: 3.14",
 ]
 dependencies = [
-  "maturin>=1.10.2",
-  "pytest>=8.4.1,<9",
-  "pytest-asyncio>=1.1.0,<2",
+  "maturin>=1.11.2",
+  "pytest>=9.0.2,<10",
+  "pytest-asyncio>=1.3.0,<2",
   "trafaret>=2.1,<3",
-  "testcontainers>=4.12.0,<5",
+  "testcontainers>=4.13.3,<5",
 ]
 
 [project.urls]
@@ -35,12 +35,12 @@ repository = "https://github.com/lablup/etcd-client-py"
 
 [project.optional-dependencies]
 dev = [
-    "ruff>=0.8.5",
-    "mypy>=1.13.0",
+    "ruff>=0.14.10",
+    "mypy>=1.19.1",
 ]
 
 [build-system]
-requires = ["maturin>=1.7,<2.0"]
+requires = ["maturin>=1.11,<2.0"]
 build-backend = "maturin"
 
 [tool.maturin]

python/etcd_client/__init__.py

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