async def __spawn_process(self) -> None:
log.debug("Spawning new process")
identity = uuid.uuid4().hex
start_future = self.__create_future()
self.__starting[identity] = start_future
process = await self.__create_process(identity)
await start_future
self.processes.add(process)
async def start(self) -> None:
self.server = await asyncio.start_server(
self.__handle_client,
sock=self.socket,
)
tasks = []
for n in range(self.workers):
log.debug("Starting worker %d", n)
tasks.append(self.__spawn_process())
await asyncio.gather(*tasks)
def __create_process(self, identity: str) -> Popen:
if self.__closing:
raise RuntimeError("Pool closed")
process = Popen([sys.executable, "-m", "aiomisc.worker_pool.process"],
stdin=PIPE,
env=os.environ)
self.__spawning[identity] = process
log.debug("Spawning new worker pool process PID: %s", process.pid)
assert process.stdin
log_level = (log.getEffectiveLevel()
if LOG_LEVEL is None else LOG_LEVEL.get())
log_format = 'color' if LOG_FORMAT is None else LOG_FORMAT.get().value
process.stdin.write(
pickle.dumps((self.address, self.__cookie, identity, log_level,
log_format)))
process.stdin.close()
return process
def _kill_process(process: Popen) -> None:
if process.returncode is not None:
return None
log.debug("Terminating worker pool process PID: %s", process.pid)
process.kill()
async def handler(start_event: asyncio.Event) -> None:
log.debug("Starting to handle client")
packet_type, salt = await receive()
assert packet_type == PacketTypes.AUTH_SALT
packet_type, digest = await receive()
assert packet_type == PacketTypes.AUTH_DIGEST
hasher = HASHER()
hasher.update(salt)
hasher.update(self.__cookie)
if digest != hasher.digest():
exc = AuthenticationError("Invalid cookie")
await send(PacketTypes.EXCEPTION, exc)
raise exc
await send(PacketTypes.AUTH_OK, True)
log.debug("Client authorized")
packet_type, identity = await receive()
assert packet_type == PacketTypes.IDENTITY
process = self.__spawning.pop(identity)
starting: asyncio.Future = self.__starting.pop(identity)
if self.initializer is not None:
initializer_done = self.__create_future()
await step(self.initializer, self.initializer_args,
dict(self.initializer_kwargs), initializer_done)
try:
await initializer_done
except Exception as e:
starting.set_exception(e)
raise
else:
starting.set_result(None)
finally:
start_event.set()
else:
starting.set_result(None)
start_event.set()
while True:
func: Callable
args: Tuple[Any, ...]
kwargs: Dict[str, Any]
result_future: asyncio.Future
process_future: asyncio.Future
(
func,
args,
kwargs,
result_future,
process_future,
) = await self.tasks.get()
try:
if process_future.done():
continue
process_future.set_result(process)
if result_future.done():
continue
await step(func, args, kwargs, result_future)
except asyncio.IncompleteReadError:
await self.__wait_process(process)
self.__on_exit(process)
result_future.set_exception(
ProcessError(
"Process {!r} exited with code {!r}".format(
process,
process.returncode,
), ), )
break
except Exception as e:
if not result_future.done():
self.loop.call_soon(result_future.set_exception, e)
if not writer.is_closing():
self.loop.call_soon(writer.close)
await self.__wait_process(process)
self.__on_exit(process)
raise