FEAT: generator support normal function (#83)

python/xoscar/api.py

@@ -16,7 +16,9 @@
 from __future__ import annotations
 
 import asyncio
+import functools
 import inspect
+import logging
 import threading
 import uuid
 from collections import defaultdict
@@ -44,6 +46,8 @@
     from .backends.config import ActorPoolConfig
     from .backends.pool import MainActorPoolType
 
+logger = logging.getLogger(__name__)
+
 
 async def create_actor(
     actor_cls: Type, *args, uid=None, address=None, **kwargs
@@ -295,6 +299,7 @@ def __init__(self, uid: str, actor_addr: str, actor_uid: str):
         self._actor_addr = actor_addr
         self._actor_uid = actor_uid
         self._actor_ref = None
+        self._gc_destroy = True
 
     async def destroy(self):
         if self._actor_ref is None:
@@ -307,15 +312,24 @@ async def destroy(self):
     def __del__(self):
         # It's not a good idea to spawn a new thread and join in __del__,
         # but currently it's the only way to GC the generator.
-        thread = threading.Thread(
-            target=asyncio.run, args=(self.destroy(),), daemon=True
-        )
-        thread.start()
-        thread.join()
+        # TODO(codingl2k1): This __del__ may hangs if the program is exiting.
+        if self._gc_destroy:
+            thread = threading.Thread(
+                target=asyncio.run, args=(self.destroy(),), daemon=True
+            )
+            thread.start()
+            thread.join()
 
     def __aiter__(self):
         return self
 
+    def __getstate__(self):
+        # Transfer gc destroy during serialization.
+        state = dict(**super().__getstate__())
+        state["_gc_destroy"] = True
+        self._gc_destroy = False
+        return state
+
     async def __anext__(self) -> T:
         if self._actor_ref is None:
             self._actor_ref = await actor_ref(
@@ -400,13 +414,7 @@ async def _async_wrapper(_gen):
             stop = object()
             try:
                 if inspect.isgenerator(gen):
-                    # to_thread is only available for Python >= 3.9
-                    if hasattr(asyncio, "to_thread"):
-                        r = await asyncio.to_thread(_wrapper, gen)
-                    else:
-                        r = await asyncio.get_event_loop().run_in_executor(
-                            None, _wrapper, gen
-                        )
+                    r = await asyncio.to_thread(_wrapper, gen)
                 elif inspect.isasyncgen(gen):
                     r = await asyncio.create_task(_async_wrapper(gen))
                 else:
@@ -415,15 +423,20 @@ async def _async_wrapper(_gen):
                         f"but a {type(gen)} is got."
                     )
             except Exception as e:
+                logger.exception(
+                    f"Destroy generator {generator_uid} due to an error encountered."
+                )
                 await self.__xoscar_destroy_generator__(generator_uid)
+                del gen  # Avoid exception hold generator reference.
                 raise e
             if r is stop:
                 await self.__xoscar_destroy_generator__(generator_uid)
+                del gen  # Avoid exception hold generator reference.
                 raise Exception("StopIteration")
             else:
                 return r
         else:
-            raise RuntimeError(f"no iterator with id: {generator_uid}")
+            raise RuntimeError(f"No iterator with id: {generator_uid}")
 
     async def __xoscar_destroy_generator__(self, generator_uid: str):
         """
@@ -434,19 +447,28 @@ async def __xoscar_destroy_generator__(self, generator_uid: str):
         generator_uid: str
             The uid of generator
         """
-        return self._generators.pop(generator_uid, None)
+        logger.debug("Destroy generator: %s", generator_uid)
+        self._generators.pop(generator_uid, None)
 
 
 def generator(func):
-    async def wrapper(obj, *args, **kwargs):
-        gen_uid = uuid.uuid1().hex
-        obj._generators[gen_uid] = func(obj, *args, **kwargs)
-        return IteratorWrapper(gen_uid, obj.address, obj.uid)
+    need_to_thread = not asyncio.iscoroutinefunction(func)
 
-    if inspect.isgeneratorfunction(func) or inspect.isasyncgenfunction(func):
-        return wrapper
-    else:
-        return func
+    @functools.wraps(func)
+    async def _wrapper(self, *args, **kwargs):
+        if need_to_thread:
+            r = await asyncio.to_thread(func, self, *args, **kwargs)
+        else:
+            r = await func(self, *args, **kwargs)
+        if inspect.isgenerator(r) or inspect.isasyncgen(r):
+            gen_uid = uuid.uuid1().hex
+            logger.debug("Create generator: %s", gen_uid)
+            self._generators[gen_uid] = r
+            return IteratorWrapper(gen_uid, self.address, self.uid)
+        else:
+            return r
+
+    return _wrapper
 
 
 class Actor(AsyncActorMixin, _Actor):

python/xoscar/core.pyx

@@ -127,7 +127,7 @@ cdef class ActorRef:
         return create_actor_ref, (self.address, self.uid)
 
     def __getattr__(self, item):
-        if item.startswith('_'):
+        if item.startswith('_') and item not in ["__xoscar_next__", "__xoscar_destroy_generator__"]:
             return object.__getattribute__(self, item)
 
         try:

python/xoscar/tests/test_generator.py

@@ -63,6 +63,32 @@ async def with_exception(self):
         raise Exception("intent raise")
         yield 2
 
+    @xo.generator
+    async def mix_gen(self, v):
+        if v == 1:
+            return self._gen()
+        elif v == 2:
+            return self._gen2()
+        else:
+            return 0
+
+    @xo.generator
+    def mix_gen2(self, v):
+        if v == 1:
+            return self._gen()
+        elif v == 2:
+            return self._gen2()
+        else:
+            return 0
+
+    def _gen(self):
+        for x in range(3):
+            yield x
+
+    async def _gen2(self):
+        for x in range(3):
+            yield x
+
     @classmethod
     def uid(cls):
         return "supervisor"
@@ -134,7 +160,19 @@ async def test_generator():
     all_gen = await superivsor_actor.get_all_generators()
     assert len(all_gen) == 0
 
-    await asyncio.create_task(superivsor_actor.with_exception())
+    r = await superivsor_actor.with_exception()
+    del r
     await asyncio.sleep(0)
     all_gen = await superivsor_actor.get_all_generators()
     assert len(all_gen) == 0
+
+    for f in [superivsor_actor.mix_gen, superivsor_actor.mix_gen2]:
+        out = []
+        async for x in await f(1):
+            out.append(x)
+        assert out == [0, 1, 2]
+        out = []
+        async for x in await f(2):
+            out.append(x)
+        assert out == [0, 1, 2]
+        assert 0 == await f(0)