Repository: mars-project/mars

Issue:
[BUG] Race condition when creating actors with parallel tasks
**Describe the bug**
When creating actors in a labeled pool using parallel asyncio tasks, allocation strategy fails because of possible race condition.

**To Reproduce**
```python
import asyncio
import os
from mars import oscar as mo
from mars.oscar.backends.allocate_strategy import IdleLabel

pool = await mo.create_actor_pool('127.0.0.1', n_process=2, labels=[None, 'numa-0', 'numa-0'])

class TestActor(mo.Actor):
    def get_pid(self):
        return os.getpid()

strategy = IdleLabel('numa-0', 'tests')
tasks = [
    mo.create_actor(TestActor, allocate_strategy=strategy, address=pool.external_address),
    mo.create_actor(TestActor, allocate_strategy=strategy, address=pool.external_address),
]
refs = await asyncio.gather(*tasks)
# outputs identical process ids, while the result should be different
print([await ref.get_pid() for ref in refs])
```

**Expected behavior**
``get_pid()`` calls return different process ids.



Discussion:

Resolving patch:
diff --git a/mars/oscar/backends/pool.py b/mars/oscar/backends/pool.py
--- a/mars/oscar/backends/pool.py
+++ b/mars/oscar/backends/pool.py
@@ -16,6 +16,7 @@
 import concurrent.futures as futures
 import contextlib
 import itertools
+import threading
 import multiprocessing
 from abc import ABC, ABCMeta, abstractmethod
 from typing import Dict, List, Type, TypeVar, Coroutine, Callable, Union, Optional
@@ -649,6 +650,7 @@ def __init__(self,
         # states
         self._allocated_actors: allocated_type = \
             {addr: dict() for addr in self._config.get_external_addresses()}
+        self._allocation_lock = threading.Lock()
 
         self.sub_processes: Dict[str, SubProcessHandle] = dict()
 
@@ -668,9 +670,12 @@ async def create_actor(self,
         with _ErrorProcessor(message_id=message.message_id,
                              protocol=message.protocol) as processor:
             allocate_strategy = message.allocate_strategy
-            # get allocated address according to corresponding strategy
-            address = allocate_strategy.get_allocated_address(
-                self._config, self._allocated_actors)
+            with self._allocation_lock:
+                # get allocated address according to corresponding strategy
+                address = allocate_strategy.get_allocated_address(
+                    self._config, self._allocated_actors)
+                # set placeholder to make sure this label is occupied
+                self._allocated_actors[address][None] = (allocate_strategy, message)
             if address == self.external_address:
                 # creating actor on main actor pool
                 result = await super().create_actor(message)
@@ -695,6 +700,9 @@ async def create_actor(self,
                         (allocate_strategy, new_create_actor_message)
                 processor.result = result
 
+            # revert placeholder
+            self._allocated_actors[address].pop(None, None)
+
         return processor.result
 
     @implements(AbstractActorPool.has_actor)


Test patch:
diff --git a/mars/oscar/backends/mars/tests/test_pool.py b/mars/oscar/backends/mars/tests/test_pool.py
--- a/mars/oscar/backends/mars/tests/test_pool.py
+++ b/mars/oscar/backends/mars/tests/test_pool.py
@@ -574,3 +574,26 @@ def is_interprocess_address(addr):
             Router.get_instance().get_internal_address(actor_ref4.address))
 
         assert await actor_ref2.add_other(actor_ref4, 3) == 13
+
+
+@pytest.mark.asyncio
+async def test_parallel_allocate_idle_label():
+    start_method = os.environ.get('POOL_START_METHOD', 'forkserver') \
+        if sys.platform != 'win32' else None
+    pool = await create_actor_pool('127.0.0.1', pool_cls=MainActorPool, n_process=2,
+                                   subprocess_start_method=start_method,
+                                   labels=[None, 'my_label', 'my_label'])
+
+    class _Actor(Actor):
+        def get_pid(self):
+            return os.getpid()
+
+    ctx = get_context()
+    strategy = IdleLabel('my_label', 'tests')
+    tasks = [
+        ctx.create_actor(_Actor, allocate_strategy=strategy, address=pool.external_address),
+        ctx.create_actor(_Actor, allocate_strategy=strategy, address=pool.external_address),
+    ]
+    refs = await asyncio.gather(*tasks)
+    # outputs identical process ids, while the result should be different
+    assert len({await ref.get_pid() for ref in refs}) == 2