Experiment: Wrap all capnp code in a context-manager to avoid segfaults (#317)

* Experiment: Wrap all capnp code in a context-manager

* Fix segfault in on_disconnect

test/test_capability.py

@@ -4,6 +4,12 @@ import capnp
 import test_capability_capnp as capability
 
 
+@pytest.fixture(autouse=True)
+async def kj_loop():
+    async with capnp.kj_loop():
+        yield
+
+
 class Server(capability.TestInterface.Server):
     def __init__(self, val=1):
         self.val = val

test/test_capability_context.py

@@ -4,6 +4,12 @@ import capnp
 import test_capability_capnp as capability
 
 
+@pytest.fixture(autouse=True)
+async def kj_loop():
+    async with capnp.kj_loop():
+        yield
+
+
 class Server(capability.TestInterface.Server):
     def __init__(self, val=1):
         self.val = val

test/test_context_manager.py

@@ -0,0 +1,241 @@
+import pytest
+import asyncio
+import socket
+
+import capnp
+import test_capability
+import test_capability_capnp as capability
+
+
+async def test_two_kj_one_asyncio():
+    async with capnp.kj_loop():
+        pass
+    async with capnp.kj_loop():
+        pass
+
+
+def test_two_kj_two_asyncio():
+    async def do():
+        async with capnp.kj_loop():
+            pass
+
+    asyncio.run(do())
+    asyncio.run(do())
+
+
+async def test_nested_kj():
+    with pytest.raises(RuntimeError) as exninfo:
+        async with capnp.kj_loop():
+            async with capnp.kj_loop():
+                pass
+    assert "The KJ event-loop is already running" in str(exninfo)
+
+
+async def test_kj_loop_leak_new_client():
+    async with capnp.kj_loop():
+        client = capability.TestInterface._new_client(test_capability.Server())
+    with pytest.raises(RuntimeError) as exninfo:
+        await client.foo(5, True)
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_kj_loop_leak_client():
+    read, write = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+        _ = capnp.TwoPartyServer(write, bootstrap=test_capability.Server())
+        client = capnp.TwoPartyClient(read)
+        cap = client.bootstrap().cast_as(capability.TestInterface)
+    with pytest.raises(RuntimeError) as exninfo:
+        await cap.foo(5, True)
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_kj_loop_leak_client2():
+    read, write = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+        _ = capnp.TwoPartyServer(write, bootstrap=test_capability.Server())
+        client = capnp.TwoPartyClient(read)
+    with pytest.raises(RuntimeError) as exninfo:
+        client.bootstrap().cast_as(capability.TestInterface)
+    assert "This client is closed" in str(exninfo)
+
+
+async def test_kj_loop_leak_client3():
+    read, write = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+        _ = capnp.TwoPartyServer(write, bootstrap=test_capability.Server())
+        client = capnp.TwoPartyClient(read).bootstrap()
+    with pytest.raises(RuntimeError) as exninfo:
+        cap = client.cast_as(capability.TestInterface)
+        await cap.foo(5, True)
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_no_kj_loop():
+    read, write = socket.socketpair()
+    with pytest.raises(RuntimeError) as exninfo:
+        await capnp.AsyncIoStream.create_connection(sock=read)
+    assert "The KJ event-loop is not running" in str(exninfo)
+    with pytest.raises(RuntimeError) as exninfo:
+        await capnp.AsyncIoStream.create_connection(sock=write)
+    assert "The KJ event-loop is not running" in str(exninfo)
+    with pytest.raises(RuntimeError) as exninfo:
+        capability.TestPipeline._new_client(test_capability.PipelineServer())
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_promise_leaking1():
+    async with capnp.kj_loop():
+        client = capability.TestInterface._new_client(test_capability.Server())
+        remote = client.foo(5, True)
+        task = asyncio.ensure_future(remote)
+        await asyncio.sleep(0)
+    with pytest.raises(capnp.KjException):
+        await task
+
+
+async def test_promise_leaking2():
+    async with capnp.kj_loop():
+        client = capability.TestInterface._new_client(test_capability.Server())
+        remote = client.foo(5, True)
+        task = asyncio.ensure_future(remote)
+    with pytest.raises(RuntimeError) as exninfo:
+        await task
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_promise_leaking3():
+    async with capnp.kj_loop():
+        client = capability.TestInterface._new_client(test_capability.Server())
+        remote = client.foo(5, True)
+    with pytest.raises(RuntimeError) as exninfo:
+        await remote
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_promise_leaking4():
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        connection = await capnp.AsyncIoStream.create_connection(sock=read)
+        client = capnp.TwoPartyClient(connection)
+        cap = client.bootstrap().cast_as(capability.TestInterface)
+        res = asyncio.ensure_future(cap.foo(5, True))
+        await asyncio.sleep(0)
+    with pytest.raises(capnp.KjException):
+        await res
+
+
+async def test_promise_leaking5():
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        connection = await capnp.AsyncIoStream.create_connection(sock=read)
+        client = capnp.TwoPartyClient(connection)
+        cap = client.bootstrap().cast_as(capability.TestInterface)
+        res = asyncio.ensure_future(cap.foo(5, True))
+    with pytest.raises(RuntimeError) as exninfo:
+        await res
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_promise_leaking6():
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        connection = await capnp.AsyncIoStream.create_connection(sock=read)
+        client = capnp.TwoPartyClient(connection)
+        cap = client.bootstrap().cast_as(capability.TestInterface)
+        res = cap.foo(5, True)
+    with pytest.raises(RuntimeError) as exninfo:
+        await res
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_kj_loop_read_message_after_close():
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+    with pytest.raises(RuntimeError) as exninfo:
+        await capability.TestSturdyRefHostId.read_async(read)
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_kj_loop_partial_read_message_after_close():
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+        message = capability.TestSturdyRefHostId.read_async(read)
+    with pytest.raises(RuntimeError) as exninfo:
+        await message
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_kj_loop_write_message_after_close():
+    _, write = socket.socketpair()
+    async with capnp.kj_loop():
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+    message = capability.TestSturdyRefHostId.new_message()
+    with pytest.raises(RuntimeError) as exninfo:
+        await message.write_async(write)
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_kj_loop_partial_write_message_after_close():
+    _, write = socket.socketpair()
+    async with capnp.kj_loop():
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+        message = capability.TestSturdyRefHostId.new_message()
+        send = message.write_async(write)
+    with pytest.raises(RuntimeError) as exninfo:
+        await send
+    assert "The KJ event-loop is not running" in str(exninfo)
+
+
+async def test_client_on_disconnect_memory():
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+        client = capnp.TwoPartyClient(read)
+    with pytest.raises(RuntimeError) as exninfo:
+        await client.on_disconnect()
+    assert "This client is closed" in str(exninfo)
+
+
+async def test_server_on_disconnect_memory():
+    _, write = socket.socketpair()
+    async with capnp.kj_loop():
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+        server = capnp.TwoPartyServer(write, bootstrap=test_capability.Server())
+    with pytest.raises(RuntimeError) as exninfo:
+        await server.on_disconnect()
+    assert "This server is closed" in str(exninfo)
+
+
+@pytest.mark.xfail(
+    strict=True,
+    reason="Fails because the promisefulfiller got destroyed. Possibly a bug in the C++ library.",
+)
+async def test_client_on_disconnect_memory2():
+    """
+    E       capnp.lib.capnp.KjException: kj/async.c++:2813: failed:
+            PromiseFulfiller was destroyed without fulfilling the promise.
+    """
+    read, _ = socket.socketpair()
+    async with capnp.kj_loop():
+        read = await capnp.AsyncIoStream.create_connection(sock=read)
+        client = capnp.TwoPartyClient(read)
+        disc = client.on_disconnect()
+    await disc
+
+
+async def test_server_on_disconnect_memory2():
+    _, write = socket.socketpair()
+    async with capnp.kj_loop():
+        write = await capnp.AsyncIoStream.create_connection(sock=write)
+        server = capnp.TwoPartyServer(write, bootstrap=test_capability.Server())
+        disc = server.on_disconnect()
+    await disc

test/test_memory_handling.py

@@ -0,0 +1,32 @@
+from types import coroutine
+import pytest
+import socket
+import gc
+
+import capnp
+import test_capability
+import test_capability_capnp as capability
+
+
+@pytest.fixture(autouse=True)
+async def kj_loop():
+    async with capnp.kj_loop():
+        yield
+
+
+@coroutine
+def wrap(p):
+    return (yield from p)
+
+
+async def test_kj_loop_await_attach():
+    read, write = socket.socketpair()
+    read = await capnp.AsyncIoStream.create_connection(sock=read)
+    write = await capnp.AsyncIoStream.create_connection(sock=write)
+    _ = capnp.TwoPartyServer(write, bootstrap=test_capability.Server())
+    client = capnp.TwoPartyClient(read).bootstrap().cast_as(capability.TestInterface)
+    t = wrap(client.foo(5, True).__await__())
+    del client
+    del read
+    gc.collect()
+    await t

test/test_response.py

@@ -1,6 +1,15 @@
+import pytest
+
+import capnp
 import test_response_capnp
 
 
+@pytest.fixture(autouse=True)
+async def kj_loop():
+    async with capnp.kj_loop():
+        yield
+
+
 class FooServer(test_response_capnp.Foo.Server):
     def __init__(self, val=1):
         self.val = val

test/test_rpc.py

@@ -9,6 +9,12 @@ import socket
 import test_capability_capnp
 
 
+@pytest.fixture(autouse=True)
+async def kj_loop():
+    async with capnp.kj_loop():
+        yield
+
+
 class Server(test_capability_capnp.TestInterface.Server):
     def __init__(self, val=100):
         self.val = val

test/test_rpc_calculator.py

@@ -2,6 +2,7 @@ import gc
 import os
 import socket
 import sys  # add examples dir to sys.path
+import pytest
 
 import capnp
 
@@ -12,6 +13,12 @@ import async_calculator_client  # noqa: E402
 import async_calculator_server  # noqa: E402
 
 
+@pytest.fixture(autouse=True)
+async def kj_loop():
+    async with capnp.kj_loop():
+        yield
+
+
 async def test_calculator():
     read, write = socket.socketpair()
     read = await capnp.AsyncIoStream.create_connection(sock=read)