84 lines
2.4 KiB
Python
84 lines
2.4 KiB
Python
from collections import defaultdict
|
|
from structio.abc import BaseIOManager, BaseKernel
|
|
from structio.core.context import Task
|
|
from structio.core.run import current_loop, current_task
|
|
import select
|
|
|
|
|
|
class SimpleIOManager(BaseIOManager):
|
|
"""
|
|
A simple, cross-platform, select()-based
|
|
I/O manager
|
|
"""
|
|
|
|
def __init__(self):
|
|
"""
|
|
Public object constructor
|
|
"""
|
|
|
|
# Maps resources to tasks
|
|
self.readers = {}
|
|
self.writers = {}
|
|
# This allows us to have a bidirectional mapping:
|
|
# we know both which tasks are using which resources
|
|
# and which resources are used by which tasks,
|
|
# without having to go through too many hoops and jumps.
|
|
self.tasks: dict[Task, list] = defaultdict(list)
|
|
|
|
def pending(self):
|
|
# We don't return bool(self.resources) because there is
|
|
# no pending I/O to do if no tasks are waiting to read or
|
|
# write, even if there's dangling resources around!
|
|
return bool(self.readers or self.writers)
|
|
|
|
def _collect_readers(self) -> list:
|
|
"""
|
|
Collects all resources that need to be read from,
|
|
so we can select() on them later
|
|
"""
|
|
|
|
result = []
|
|
for resource in self.readers:
|
|
result.append(resource)
|
|
return result
|
|
|
|
def _collect_writers(self) -> list:
|
|
"""
|
|
Collects all resources that need to be written to,
|
|
so we can select() on them later
|
|
"""
|
|
|
|
result = []
|
|
for resource in self.writers:
|
|
result.append(resource)
|
|
return result
|
|
|
|
def wait_io(self):
|
|
kernel: BaseKernel = current_loop()
|
|
readable, writable, _ = select.select(
|
|
self._collect_readers(),
|
|
self._collect_writers(),
|
|
[],
|
|
kernel.get_closest_deadline(),
|
|
)
|
|
for read_ready in readable:
|
|
kernel.reschedule(self.readers[read_ready])
|
|
for write_ready in writable:
|
|
kernel.reschedule(self.writers[write_ready])
|
|
|
|
def request_read(self, rsc):
|
|
task = current_task()
|
|
self.readers[rsc] = task
|
|
|
|
def request_write(self, rsc):
|
|
task = current_task()
|
|
self.writers[rsc] = task
|
|
|
|
def release(self, resource):
|
|
self.readers.pop(resource, None)
|
|
self.writers.pop(resource, None)
|
|
|
|
def release_task(self, task: Task):
|
|
for resource in self.tasks[task]:
|
|
self.release(resource)
|