class Pool(object):
def __init__(self, size=1000):
self._size = size
self._lock = Semaphore(size)
self._running_jobs = 0
self._end_event = Event()
self._end_event.set()
def spawn(self, func, *args, **kw):
self._lock.acquire()
self._running_jobs += 1
self._end_event.clear()
return evergreen.spawn(self._runner, func, args, kw)
def join(self, timeout=None):
return self._end_event.wait(timeout)
def _runner(self, func, args, kw):
try:
func(*args, **kw)
finally:
self._running_jobs -= 1
if self._running_jobs == 0:
self._end_event.set()
self._lock.release()
def __repr__(self):
return '<%s(size=%d), %d running jobs>' % (
self.__class__.__name__, self._size, self._running_jobs)
class Pool(object):
def __init__(self, size=1000):
self._size = size
self._lock = Semaphore(size)
self._running_jobs = 0
self._end_event = Event()
self._end_event.set()
def spawn(self, func, *args, **kw):
self._lock.acquire()
self._running_jobs += 1
self._end_event.clear()
return evergreen.spawn(self._runner, func, args, kw)
def join(self, timeout=None):
return self._end_event.wait(timeout)
def _runner(self, func, args, kw):
try:
func(*args, **kw)
finally:
self._running_jobs -= 1
if self._running_jobs == 0:
self._end_event.set()
self._lock.release()
def __repr__(self):
return '<%s(size=%d), %d running jobs>' % (self.__class__.__name__, self._size, self._running_jobs)
class IOHandler(object):
def __init__(self, fd):
self.fd = fd
self._read_closed = False
self._write_closed = False
self._read_event = Event()
self._write_event = Event()
def wait_read(self, timeout=None, timeout_exc=None):
if self._read_closed:
raise cancel_wait_ex
self._read_event.clear()
loop = evergreen.current.loop
handler = loop.add_reader(self.fd, self._read_event.set)
try:
self._wait(self._read_event, timeout, timeout_exc)
if self._read_closed:
raise cancel_wait_ex
finally:
handler.cancel()
loop.remove_reader(self.fd)
def wait_write(self, timeout=None, timeout_exc=None):
if self._write_closed:
raise cancel_wait_ex
self._write_event.clear()
loop = evergreen.current.loop
handler = loop.add_writer(self.fd, self._write_event.set)
try:
self._wait(self._write_event, timeout, timeout_exc)
if self._write_closed:
raise cancel_wait_ex
finally:
handler.cancel()
loop.remove_writer(self.fd)
def close(self, read=True, write=True):
if read:
self._read_closed = True
self._read_event.set()
if write:
self._write_closed = True
self._write_event.set()
def _wait(self, event, timeout, timeout_exc):
r = event.wait(timeout)
if not r and timeout_exc:
raise timeout_exc
def __repr__(self):
return '<%s fd=%d>' % (self.__class__.__name__, self.fd)
class Result(object):
def __init__(self):
self._event = Event()
self._result = None
self._exception = None
self._used = False
def is_set(self):
return self._event.is_set()
def set_result(self, value):
if self.is_set():
raise RuntimeError('already set')
self._result = value
self._event.set()
def set_exception(self, exc):
if self.is_set():
raise RuntimeError('already set')
self._exception = exc
self._event.set()
def wait(self, timeout=None):
if self._used:
raise RuntimeError('already used, clear it in order to use it again')
self._event.wait(timeout)
try:
if self.is_set():
return self._get_result()
return None
finally:
self._result = self._exception = None
def clear(self):
self._event.clear()
self._result = None
self._exception = None
self._used = False
def _get_result(self):
self._used = True
if self._exception is not None:
six.reraise(type(self._exception), self._exception)
else:
return self._result
class Result(object):
def __init__(self):
self._event = Event()
self._result = None
self._exception = None
self._used = False
def is_set(self):
return self._event.is_set()
def set_result(self, value):
if self.is_set():
raise RuntimeError('already set')
self._result = value
self._event.set()
def set_exception(self, exc):
if self.is_set():
raise RuntimeError('already set')
self._exception = exc
self._event.set()
def wait(self, timeout=None):
if self._used:
raise RuntimeError(
'already used, clear it in order to use it again')
self._event.wait(timeout)
try:
if self.is_set():
return self._get_result()
return None
finally:
self._result = self._exception = None
def clear(self):
self._event.clear()
self._result = None
self._exception = None
self._used = False
def _get_result(self):
self._used = True
if self._exception is not None:
six.reraise(type(self._exception), self._exception)
else:
return self._result
class Channel(object):
def __init__(self):
self._send_lock = Lock()
self._recv_lock = Lock()
self._new_data = Event()
self._recv_data = Event()
self._data = None
def send(self, data):
with self._send_lock:
self._data = data
self._new_data.set()
self._recv_data.wait()
self._recv_data.clear()
def send_exception(self, exc_type, exc_value=None, exc_tb=None):
self.send(_Bomb(exc_type, exc_value, exc_tb))
def receive(self):
with self._recv_lock:
self._new_data.wait()
data, self._data = self._data, None
self._new_data.clear()
self._recv_data.set()
if isinstance(data, _Bomb):
data.raise_()
else:
return data
def __iter__(self):
return self
def next(self):
return self.receive()
if six.PY3:
__next__ = next
del next
class UDPEndpoint(object):
def __init__(self):
loop = evergreen.current.loop
self._handle = pyuv.UDP(loop._loop)
self._closed = False
self._receive_result = Result()
self._pending_writes = 0
self._flush_event = Event()
self._flush_event.set()
self._sockname = None
@property
def sockname(self):
self._check_closed()
if self._sockname is None:
self._sockname = self._handle.getsockname()
return self._sockname
def bind(self, addr):
self._check_closed()
self._handle.bind(addr)
def send(self, data, addr):
self._check_closed()
self._handle.send(addr, data, self.__send_cb)
if self._pending_writes == 0:
self._flush_event.clear()
self._pending_writes += 1
def receive(self):
self._check_closed()
with self._receive_result:
self._handle.start_recv(self.__receive_cb)
return self._receive_result.get()
def flush(self):
self._check_closed()
self._flush_event.wait()
def close(self):
if self._closed:
return
self._closed = True
self._handle.close()
def _check_closed(self):
if self._closed:
raise UDPError("endpoint is closed")
def __send_cb(self, handle, error):
self._pending_writes -= 1
if self._pending_writes == 0:
self._flush_event.set()
if error is not None:
log.debug("send failed: %d %s", error, pyuv.errno.strerror(error))
evergreen.current.loop.call_soon(self.close)
def __receive_cb(self, handle, addr, flags, data, error):
self._handle.stop_recv()
if error is not None:
self._receive_result.set_exception(UDPError(error, errno.strerror(error)))
else:
self._receive_result.set_value((data, addr))
class BaseStream(AbstractBaseStream):
"""Base class for streams implemented using pyuv Stream objects as the underlying mechanism
"""
def __init__(self):
super(BaseStream, self).__init__()
self._read_result = Result()
self._shutdown_result = Result()
self._pending_writes = 0
self._flush_event = Event()
self._flush_event.set()
def flush(self):
self._check_closed()
self._flush_event.wait()
def _read(self, n):
with self._read_result:
try:
self._handle.start_read(self.__read_cb)
except self.error_cls:
self.close()
raise
try:
data = self._read_result.get()
except self.error_cls as e:
self.close()
if e.args[0] != errno.EOF:
raise
else:
self._read_buffer.feed(data)
def _write(self, data):
try:
self._handle.write(data, self.__write_cb)
except self.error_cls:
self.close()
raise
if self._pending_writes == 0:
self._flush_event.clear()
self._pending_writes += 1
return self._handle.write_queue_size == 0
def _shutdown(self):
with self._shutdown_result:
self._handle.shutdown(self.__shutdown_cb)
self._shutdown_result.get()
def _close(self):
self._handle.close()
def __read_cb(self, handle, data, error):
self._handle.stop_read()
if error is not None:
self._read_result.set_exception(self.error_cls(error, pyuv.errno.strerror(error)))
else:
self._read_result.set_value(data)
def __write_cb(self, handle, error):
self._pending_writes -= 1
if self._pending_writes == 0:
self._flush_event.set()
if error is not None:
log.debug('write failed: %d %s', error, pyuv.errno.strerror(error))
evergreen.current.loop.call_soon(self.close)
def __shutdown_cb(self, handle, error):
if error is not None:
self._shutdown_result.set_exception(self.error_cls(error, pyuv.errno.strerror(error)))
else:
self._shutdown_result.set_value(None)
class BaseStream(AbstractBaseStream):
"""Base class for streams implemented using pyuv Stream objects as the underlying mechanism
"""
def __init__(self):
super(BaseStream, self).__init__()
self._read_result = Result()
self._shutdown_result = Result()
self._pending_writes = 0
self._flush_event = Event()
self._flush_event.set()
def flush(self):
self._check_closed()
self._flush_event.wait()
def _read(self, n):
with self._read_result:
try:
self._handle.start_read(self.__read_cb)
except self.error_cls:
self.close()
raise
try:
data = self._read_result.get()
except self.error_cls as e:
self.close()
if e.args[0] != errno.EOF:
raise
else:
self._read_buffer.feed(data)
def _write(self, data):
try:
self._handle.write(data, self.__write_cb)
except self.error_cls:
self.close()
raise
if self._pending_writes == 0:
self._flush_event.clear()
self._pending_writes += 1
return self._handle.write_queue_size == 0
def _shutdown(self):
with self._shutdown_result:
self._handle.shutdown(self.__shutdown_cb)
self._shutdown_result.get()
def _close(self):
self._handle.close()
def __read_cb(self, handle, data, error):
self._handle.stop_read()
if error is not None:
self._read_result.set_exception(
self.error_cls(error, pyuv.errno.strerror(error)))
else:
self._read_result.set_value(data)
def __write_cb(self, handle, error):
self._pending_writes -= 1
if self._pending_writes == 0:
self._flush_event.set()
if error is not None:
log.debug('write failed: %d %s', error, pyuv.errno.strerror(error))
evergreen.current.loop.call_soon(self.close)
def __shutdown_cb(self, handle, error):
if error is not None:
self._shutdown_result.set_exception(
self.error_cls(error, pyuv.errno.strerror(error)))
else:
self._shutdown_result.set_value(None)
class UDPEndpoint(object):
def __init__(self):
loop = evergreen.current.loop
self._handle = pyuv.UDP(loop._loop)
self._closed = False
self._receive_result = Result()
self._pending_writes = 0
self._flush_event = Event()
self._flush_event.set()
self._sockname = None
@property
def sockname(self):
self._check_closed()
if self._sockname is None:
self._sockname = self._handle.getsockname()
return self._sockname
def bind(self, addr):
self._check_closed()
self._handle.bind(addr)
def send(self, data, addr):
self._check_closed()
self._handle.send(addr, data, self.__send_cb)
if self._pending_writes == 0:
self._flush_event.clear()
self._pending_writes += 1
def receive(self):
self._check_closed()
with self._receive_result:
self._handle.start_recv(self.__receive_cb)
return self._receive_result.get()
def flush(self):
self._check_closed()
self._flush_event.wait()
def close(self):
if self._closed:
return
self._closed = True
self._handle.close()
def _check_closed(self):
if self._closed:
raise UDPError('endpoint is closed')
def __send_cb(self, handle, error):
self._pending_writes -= 1
if self._pending_writes == 0:
self._flush_event.set()
if error is not None:
log.debug('send failed: %d %s', error, pyuv.errno.strerror(error))
evergreen.current.loop.call_soon(self.close)
def __receive_cb(self, handle, addr, flags, data, error):
self._handle.stop_recv()
if error is not None:
self._receive_result.set_exception(
UDPError(error, errno.strerror(error)))
else:
self._receive_result.set_value((data, addr))
