DefStream 0.1.13__44-py3-none-any.whl

DefStream/Worker/Startup.py

@@ -0,0 +1,174 @@
+from collections.abc import Mapping
+from os import environ, getpid
+import sys
+from typing import (
+	Any,
+	TYPE_CHECKING,
+)
+
+if TYPE_CHECKING:
+	from .Worker import Worker
+
+__all__ = (
+	'startup',
+)
+
+
+LOGO = r"""
+▄▄▄▄▄▄           ▄▄  ▄▄▄▄▄▄▄                                  
+███▀▀██▄        ██  █████▀▀▀  ██                              
+███  ███ ▄█▀█▄ ▀██▀  ▀████▄  ▀██▀▀ ████▄ ▄█▀█▄  ▀▀█▄ ███▄███▄ 
+███  ███ ██▄█▀  ██     ▀████  ██   ██ ▀▀ ██▄█▀ ▄█▀██ ██ ██ ██ 
+██████▀  ▀█▄▄▄  ██  ███████▀  ██   ██    ▀█▄▄▄ ▀█▄██ ██ ██ ██ 
+
+Just def. Let workers do the work. You don't need to await this.
+"""
+
+
+def supports_color() -> bool:
+	return (
+		getattr(sys.stdout, 'isatty', lambda: False)()
+		and not environ.get('NO_COLOR')
+	)
+
+
+def colorize_logo(logo: str) -> str:
+	colors = (
+		'\033[38;5;39m',
+		'\033[38;5;45m',
+		'\033[38;5;51m',
+		'\033[38;5;45m',
+		'\033[38;5;39m',
+	)
+	reset = '\033[0m'
+	colorized_lines = []
+	art_line = 0
+	for line in logo.splitlines(keepends=True):
+		if line.strip() and art_line < len(colors):
+			newline = '\n' if line.endswith('\n') else ''
+			colorized_lines.append(
+				f'{colors[art_line]}{line[:-1] if newline else line}'
+				f'{reset}{newline}',
+			)
+			art_line += 1
+		else:
+			colorized_lines.append(line)
+	return ''.join(colorized_lines)
+
+
+def subscription_value(
+	parameters: Any,
+	name: str,
+) -> Any:
+	if isinstance(
+		parameters,
+		str,
+	):
+		return parameters if name == 'queue' else None
+	if isinstance(
+		parameters,
+		Mapping,
+	):
+		return parameters.get(name)
+	return getattr(
+		parameters,
+		name,
+		None,
+	)
+
+
+def startup(worker: 'Worker') -> None:
+	broker = type(worker.broker)
+	broker_name = broker.__qualname__
+	for base in broker.__mro__:
+		if '.Broker.Implements.' in base.__module__:
+			broker_name = base.__module__.split('.Broker.Implements.')[1].split('.')[0]
+			break
+	pool = type(worker.pool)
+	conf = getattr(
+		worker.broker,
+		'conf',
+		None,
+	)
+	host = getattr(
+		conf,
+		'host',
+		None,
+	)
+	port = getattr(
+		conf,
+		'port',
+		None,
+	)
+	if isinstance(
+		host,
+		str,
+	):
+		address = f'[{host}]' if ':' in host and not host.startswith('[') else host
+		broker_name += f'://{address}'
+		if isinstance(
+			port,
+			(
+				str,
+				int,
+			),
+		):
+			broker_name += f':{port}'
+		vhost = getattr(
+			conf,
+			'vhost',
+			None,
+		)
+		if isinstance(
+			vhost,
+			str,
+		) and vhost:
+			broker_name += '/' + vhost.lstrip('/')
+	logo = LOGO.rstrip()
+	starting_worker = 'Starting Worker'
+	if supports_color():
+		logo = colorize_logo(logo)
+		starting_worker = f'\033[1;38;5;51m{starting_worker}\033[0m'
+	lines = [
+		logo,
+		'',
+		starting_worker,
+		f'  PID: {getpid()}',
+		f'  Broker: {broker_name}',
+	]
+	lines.extend((
+		f'  Pool: {pool.__name__} (concurrency={worker.pool.size()})',
+		f'  Tasks: {len(worker.tasks)} tasks',
+	))
+	queues = []
+	for task in worker.tasks:
+		name = task.__qualname__
+		queue = subscription_value(
+			task.parameters,
+			'queue',
+		)
+		lines.append(f'    - {name}')
+		lines.append(f'      queue={queue!r}' if queue is not None else '      queue=<adapter-defined>')
+		if queue is not None and queue not in queues:
+			queues.append(queue)
+		for option in (
+			'group',
+			'qos',
+			'offset',
+			'error',
+		):
+			value = subscription_value(
+				task.parameters,
+				option,
+			)
+			if value is not None:
+				lines.append(f'      {option}={value!r}')
+		policy = type(task.on_error).__name__ if task.on_error is not None else 'broker default'
+		lines.append(f'      on_error={policy}')
+	lines.append(f'  Queues: {len(queues)} queues')
+	lines.extend(f'    - {queue!r}' for queue in queues)
+	lines.append('\n')
+	print(
+		'\n'.join(lines),
+		flush=True,
+	)

DefStream/Worker/Worker.py

@@ -0,0 +1,421 @@
+from abc import (
+	ABC,
+	abstractmethod,
+)
+from collections.abc import Mapping
+from logging import getLogger
+from threading import Event as PyEvent
+from typing import (
+	Any,
+	List,
+	Callable,
+	Optional,
+	Sequence,
+	Type,
+	Tuple,
+	Union,
+)
+
+from ..Broker import (
+	Connection,
+	GetConsumer,
+	Consumer,
+	EventHandler,
+)
+from ..Broker.Event import Event
+from .Message import task_parameters
+from ..Inspector import Inspector
+from .Pool import (
+	Pool,
+	PoolFull,
+)
+from .Run import Run
+from .Executor import run_cleanups
+from .Startup import startup
+
+__all__ = (
+	'Worker',
+	'WorkerSetup',
+	'WorkerCleanup',
+)
+
+class WorkerSetup(ABC):
+	@abstractmethod
+	def __call__(
+		self,
+		worker: 'Worker',
+	) -> None:
+		raise NotImplementedError
+
+
+class WorkerCleanup(ABC):
+	@abstractmethod
+	def __call__(
+		self,
+		worker: 'Worker',
+	) -> None:
+		raise NotImplementedError
+
+
+def worker_hooks(
+	hooks: Optional[Union[Callable[['Worker'], None], Sequence[Callable[['Worker'], None]]]],
+) -> Tuple[Callable[['Worker'], None], ...]:
+	result = () if hooks is None else tuple(hooks) if isinstance(
+		hooks,
+		Sequence,
+	) else (hooks,)
+	if not all(callable(hook) for hook in result):
+		raise TypeError('Worker hooks must be callable')
+	return result
+
+
+class Worker:
+	logger = getLogger('DefStream.Worker')
+
+	def __init__(
+		self,
+		broker: Connection,
+		pool: Pool,
+		setups: Optional[Union[WorkerSetup, Sequence[WorkerSetup]]] = None,
+		cleanups: Optional[Union[WorkerCleanup, Sequence[WorkerCleanup]]] = None,
+		inspector: Optional[Inspector] = None,
+	):
+		self.broker = broker
+		self.pool = pool
+		self.tasks: List[Type[Run]] = []
+		self.consumer: Optional[Consumer] = None
+		self.started = False
+		self.running = PyEvent()
+		self.stopping = PyEvent()
+		self.finished = PyEvent()
+		self.pending = 0
+		self.setups = worker_hooks(setups)
+		self.cleanups = worker_hooks(cleanups)
+		self.inspector = inspector
+
+	def add(
+		self,
+		run: Type[Run],
+	) -> None:
+		if self.started:
+			raise RuntimeError('Cannot register runs after starting the worker')
+		task_parameters(run)
+		if run in self.tasks:
+			raise ValueError('Run is already registered')
+		if any(
+			registered.name == run.name
+			for registered in self.tasks
+		):
+			raise ValueError('Duplicate Run name: ' + run.name)
+		self.tasks.append(run)
+		self.logger.debug(
+			'Task registered: task=%s.%s',
+			run.__module__,
+			run.__qualname__,
+		)
+
+	def run(self) -> None:
+		if self.started:
+			raise RuntimeError('Worker can only be started once')
+		if not self.tasks:
+			raise RuntimeError('Register at least one Run before starting')
+		self.started = True
+		self.logger.info(
+			'Starting Worker: tasks=%s',
+			len(self.tasks),
+		)
+		try:
+			startup(self)
+			if self.inspector is not None:
+				self.inspector.start(self)
+			self.logger.debug('Connecting to broker')
+			self.broker.connect()
+			self.logger.debug('Broker connected')
+			for setup in self.setups:
+				self.logger.debug(
+					'Starting Worker setup: hook=%s',
+					type(setup).__qualname__,
+				)
+				setup(self)
+				self.logger.debug(
+					'Worker setup completed: hook=%s',
+					type(setup).__qualname__,
+				)
+			if not isinstance(self.broker, GetConsumer):
+				raise TypeError('Worker broker must support consumer creation')
+			self.consumer = self.broker.consumer()
+			self.logger.debug('Consumer created')
+			for run in self.tasks:
+				self.logger.debug(
+					'Subscribing task: task=%s.%s',
+					run.__module__,
+					run.__qualname__,
+				)
+				self.consumer.subs(
+					run.parameters,
+					Handler(
+						self,
+						run,
+					),
+				)
+				self.logger.debug(
+					'Task subscribed: task=%s.%s',
+					run.__module__,
+					run.__qualname__,
+				)
+			if not self.stopping.is_set():
+				self.running.set()
+				self.logger.info('Worker ready to consume')
+				self.consumer.run()
+		finally:
+			self.running.clear()
+			self.finished.set()
+			if self.pending:
+				self.logger.error(
+					'Consumer exited with unconfirmed deliveries: pending=%s',
+					self.pending,
+				)
+			try:
+				self.logger.info(
+					'Draining Worker tasks: pending=%s',
+					self.pending,
+				)
+				self.pool.stop()
+			finally:
+				try:
+					self.cleanup()
+				finally:
+					try:
+						if self.consumer is not None:
+							self.consumer.close()
+					finally:
+						try:
+							self.broker.close()
+						finally:
+							if self.inspector is not None and self.inspector.worker is self:
+								self.inspector.stop()
+			self.logger.info('Worker shutdown completed')
+
+	def cleanup(self) -> None:
+		self.logger.debug(
+			'Starting Worker cleanup: hooks=%s',
+			len(self.cleanups),
+		)
+		run_cleanups(
+			self.cleanups,
+			self,
+		)
+		self.logger.debug('Worker cleanup completed')
+
+	def stop(self) -> None:
+		if not self.stopping.is_set():
+			self.logger.info(
+				'Worker shutdown requested: pending=%s',
+				self.pending,
+			)
+		self.stopping.set()
+		if self.consumer is not None and not self.finished.is_set():
+			self.consumer.call(self.try_stop)
+
+	def try_stop(self) -> None:
+		if self.stopping.is_set() and self.pending == 0 and self.consumer is not None:
+			self.logger.debug('No pending tasks; requesting Consumer stop')
+			self.consumer.stop()
+
+	def receive(
+		self,
+		run: Type[Run],
+		event: Event,
+	) -> None:
+		if self.stopping.is_set():
+			self.logger.debug(
+				'Requeuing message received during shutdown: task=%s.%s',
+				run.__module__,
+				run.__qualname__,
+			)
+			event.nack()
+			return
+		if not isinstance(
+			event.body,
+			Mapping,
+		):
+			self.failed(
+				run,
+				event,
+				TypeError('Run message body must be a mapping'),
+			)
+			return
+		self.pending += 1
+		consumer = self.consumer
+		if consumer is None:
+			self.pending -= 1
+			raise RuntimeError('Worker consumer is not initialized')
+		inspector = self.inspector
+		task = run.name
+		self.logger.debug(
+			'Task message received: task=%s pending=%s',
+			task,
+			self.pending,
+		)
+		contexts = inspector.begin(task) if inspector is not None else ()
+
+		def schedule(
+			error: Optional[BaseException] = None,
+			result: Any = None,
+		) -> None:
+			if inspector is not None:
+				inspector.end(
+					task,
+					contexts,
+					error is None,
+				)
+			if self.finished.is_set():
+				self.logger.error(
+					'Cannot confirm delivery after Consumer exit: task=%s',
+					task,
+				)
+				return
+			self.logger.debug(
+				'Scheduling task result callback: task=%s success=%s',
+				task,
+				error is None,
+			)
+			consumer.call(lambda: self.complete(
+				run,
+				event,
+				error,
+				result,
+			))
+
+		try:
+			self.logger.debug(
+				'Submitting task to Pool: task=%s',
+				task,
+			)
+			self.pool.submit(
+				run,
+				dict(event.body),
+				callback=lambda result: schedule(result=result),
+				error_callback=schedule,
+			)
+			self.logger.debug(
+				'Pool task submission accepted: task=%s',
+				task,
+			)
+		except PoolFull:
+			self.logger.debug(
+				'Requeuing message because Pool is full: task=%s',
+				task,
+			)
+			if inspector is not None:
+				inspector.end(
+					task,
+					contexts,
+					None,
+				)
+			self.pending -= 1
+			event.nack()
+			return
+		except Exception:
+			if inspector is not None:
+				inspector.end(
+					task,
+					contexts,
+					None,
+				)
+			self.pending -= 1
+			raise
+
+	def complete(
+		self,
+		run: Type[Run],
+		event: Event,
+		error: Optional[BaseException] = None,
+		result: Any = None,
+	) -> None:
+		try:
+			if error is not None:
+				self.failed(
+					run,
+					event,
+					error,
+				)
+			else:
+				if self.consumer is None:
+					raise RuntimeError('Worker consumer is not initialized')
+				self.consumer.complete(
+					event,
+					result,
+				)
+				self.logger.debug(
+					'Consumer completion returned: task=%s.%s',
+					run.__module__,
+					run.__qualname__,
+				)
+		finally:
+			self.pending -= 1
+			self.logger.debug(
+				'Task completion handling finished: pending=%s',
+				self.pending,
+			)
+			self.try_stop()
+
+	def failed(
+		self,
+		run: Type[Run],
+		event: Event,
+		error: BaseException,
+	) -> None:
+		if isinstance(
+			error,
+			TypeError,
+		) and not isinstance(
+			event.body,
+			Mapping,
+		):
+			self.logger.warning(
+				'Rejected a non-mapping task message: task=%s.%s',
+				run.__module__,
+				run.__qualname__,
+			)
+		else:
+			self.logger.error(
+				'Task failed: task=%s.%s',
+				run.__module__,
+				run.__qualname__,
+				exc_info=(
+					type(error),
+					error,
+					error.__traceback__,
+				),
+			)
+		if self.consumer is None:
+			raise RuntimeError('Worker consumer is not initialized')
+		self.consumer.fail(
+			event,
+			run.on_error,
+		)
+		self.logger.debug(
+			'Consumer failure policy returned: task=%s.%s',
+			run.__module__,
+			run.__qualname__,
+		)
+
+
+class Handler(EventHandler):
+	def __init__(
+		self,
+		worker: Worker,
+		run: Type[Run],
+	):
+		self.worker = worker
+		self.run = run
+
+	def __call__(
+		self,
+		event: Event,
+	) -> None:
+		self.worker.receive(
+			self.run,
+			event,
+		)

DefStream/Worker/__init__.py

@@ -0,0 +1,47 @@
+from .Worker import (
+	Worker,
+	WorkerSetup,
+	WorkerCleanup,
+)
+from .Message import (
+	declare,
+	emit,
+	request,
+)
+from .Run import (
+	Run,
+	Context,
+	RunSetup,
+	RunComplete,
+	RunError,
+)
+from .Specification import Specification
+from .Pool import (
+	Pool,
+	PoolFull,
+	PoolSetup,
+	PoolCleanup,
+	ThreadPool,
+	ProcessPool,
+)
+
+__all__ = (
+	'Worker',
+	'declare',
+	'emit',
+	'request',
+	'WorkerSetup',
+	'WorkerCleanup',
+	'Run',
+	'Specification',
+	'Pool',
+	'PoolFull',
+	'PoolSetup',
+	'PoolCleanup',
+	'ThreadPool',
+	'ProcessPool',
+	'Context',
+	'RunSetup',
+	'RunComplete',
+	'RunError',
+)