Repository: ray-project/ray

Issue:
[Asyncio] InvalidStateError when multiple awaits on same oid
<!--Please include [tune], [rllib], [autoscaler] etc. in the issue title if relevant-->

### What is the problem?

*Ray version and other system information (Python version, TensorFlow version, OS):*

### Reproduction (REQUIRED)
Please provide a script that can be run to reproduce the issue. The script should have **no external library dependencies** (i.e., use fake or mock data / environments):

If we cannot run your script, we cannot fix your issue.
```python
import ray
import time

ray.init()
@ray.remote
def f():
    time.sleep(5)

oid = f.remote()
await asyncio.wait_for(oid, timeout=1)
await asyncio.wait_for(oid, timeout=1)
```

Output
```
Exception in callback get_async.<locals>.done_callback(<Future finis... result=None)>) at /Users/simonmo/Desktop/ray/ray/python/ray/async_compat.py:65
handle: <Handle get_async.<locals>.done_callback(<Future finis... result=None)>) at /Users/simonmo/Desktop/ray/ray/python/ray/async_compat.py:65>
Traceback (most recent call last):
  File "/Users/simonmo/miniconda3/lib/python3.6/asyncio/events.py", line 145, in _run
    self._callback(*self._args)
  File "/Users/simonmo/Desktop/ray/ray/python/ray/async_compat.py", line 83, in done_callback
    user_future.set_result(result.result)
asyncio.base_futures.InvalidStateError: invalid state
```
- [ ] I have verified my script runs in a clean environment and reproduces the issue.
- [ ] I have verified the issue also occurs with the [latest wheels](https://docs.ray.io/en/latest/installation.html).



Discussion:
Oh nice. This must be related to @allenyin55 's issue right? 

Resolving patch:
diff --git a/python/ray/async_compat.py b/python/ray/async_compat.py
--- a/python/ray/async_compat.py
+++ b/python/ray/async_compat.py
@@ -75,6 +75,11 @@ def done_callback(future):
             # Result from direct call.
             assert isinstance(result, AsyncGetResponse), result
             if result.plasma_fallback_id is None:
+                # If this future has result set already, we just need to
+                # skip the set result/exception procedure.
+                if user_future.done():
+                    return
+
                 if isinstance(result.result, ray.exceptions.RayTaskError):
                     ray.worker.last_task_error_raise_time = time.time()
                     user_future.set_exception(


Test patch:
diff --git a/python/ray/tests/test_asyncio.py b/python/ray/tests/test_asyncio.py
--- a/python/ray/tests/test_asyncio.py
+++ b/python/ray/tests/test_asyncio.py
@@ -5,6 +5,7 @@
 import sys
 
 import ray
+from ray.test_utils import SignalActor
 
 
 def test_asyncio_actor(ray_start_regular_shared):
@@ -177,6 +178,26 @@ async def plasma_get(self, plasma_object):
     assert ray.get(getter.plasma_get.remote([plasma_object])) == 2
 
 
+@pytest.mark.asyncio
+async def test_asyncio_double_await(ray_start_regular_shared):
+    # This is a regression test for
+    # https://github.com/ray-project/ray/issues/8841
+
+    signal = SignalActor.remote()
+    waiting = signal.wait.remote()
+
+    future = waiting.as_future()
+    with pytest.raises(asyncio.TimeoutError):
+        await asyncio.wait_for(future, timeout=0.1)
+    assert future.cancelled()
+
+    # We are explicitly waiting multiple times here to test asyncio state
+    # override.
+    await signal.send.remote()
+    await waiting
+    await waiting
+
+
 if __name__ == "__main__":
     import pytest
     sys.exit(pytest.main(["-v", __file__]))