fix: Update the usage of timeouts

lablup · Sep 6, 2023 · 8a23ea8 · 8a23ea8
1 parent 14305f4
commit 8a23ea8
Show file tree

Hide file tree

Showing 3 changed files with 8 additions and 14 deletions.
diff --git a/examples/simple-client-thrift.py b/examples/simple-client-thrift.py
@@ -5,14 +5,11 @@
 import sys
 import textwrap
 
-from async_timeout import timeout
-
-from callosum.rpc import Peer, RPCUserError
-from callosum.serialize import noop_serializer, noop_deserializer
+import thriftpy2 as thriftpy
 from callosum.lower.zeromq import ZeroMQAddress, ZeroMQRPCTransport
+from callosum.rpc import Peer, RPCUserError
+from callosum.serialize import noop_deserializer, noop_serializer
 from callosum.upper.thrift import ThriftClientAdaptor
-import thriftpy2 as thriftpy
-
 
 simple_thrift = thriftpy.load(
     str(pathlib.Path(__file__).parent / "simple.thrift"), module_name="simple_thrift"
@@ -47,7 +44,7 @@ async def call() -> None:
             print(textwrap.indent(e.traceback, prefix="| "))
 
         try:
-            with timeout(0.5):
+            async with asyncio.timeout(0.5):
                 await peer.invoke("simple", adaptor.long_delay())
         except asyncio.TimeoutError:
             print(

diff --git a/examples/simple-client.py b/examples/simple-client.py
@@ -1,16 +1,14 @@
 import asyncio
 import json
-import random
 import os
+import random
 import secrets
 import sys
 import textwrap
 import traceback
 
-from async_timeout import timeout
-
-from callosum.rpc import Peer, RPCUserError
 from callosum.lower.zeromq import ZeroMQAddress, ZeroMQRPCTransport
+from callosum.rpc import Peer, RPCUserError
 
 
 async def test_simple(peer, initial_delay: float = 0):
@@ -41,7 +39,7 @@ async def test_simple(peer, initial_delay: float = 0):
 
 async def test_timeout(peer):
     try:
-        with timeout(0.5):
+        async with asyncio.timeout(0.5):
             await peer.invoke(
                 "long_delay",
                 {

diff --git a/src/callosum/rpc/channel.py b/src/callosum/rpc/channel.py
@@ -18,7 +18,6 @@
 
 import attrs
 from aiotools import aclosing
-from async_timeout import timeout
 
 from ..abc import (
     AbstractChannel,
@@ -327,7 +326,7 @@ async def invoke(
         try:
             request: RPCMessage
             server_cancelled = False
-            with timeout(invoke_timeout):
+            async with asyncio.timeout(invoke_timeout):
                 if callable(body):
                     # The user is using an upper-layer adaptor.
                     async with aclosing(body()) as agen: