You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.

189 lines
5.4 KiB

# Copyright 2019 Red Hat
# Licensed under the Apache License, Version 2.0 (the "License"); you may
# not use this file except in compliance with the License. You may obtain
# a copy of the License at
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations
# under the License.
import time
import logging
import queue
import threading
from nodepool.driver import Provider
class Task:
"""Base task class for use with :py:class:`TaskManager`
Subclass this to implement your own tasks.
Set the `name` field to the name of your task and override the
:py:meth:`main` method.
Keyword arguments to the constructor are stored on `self.args` for
use by the :py:meth:`main` method.
name = "task_name"
def __init__(self, **kw):
self._wait_event = threading.Event()
self._exception = None
self._traceback = None
self._result = None
self.args = kw
def done(self, result):
self._result = result
def exception(self, e):
self._exception = e
def wait(self):
"""Call this method after submitting the task to the TaskManager to
receieve the results."""
if self._exception:
raise self._exception
return self._result
def run(self, manager):
except Exception as e:
def main(self, manager):
"""Implement the work of the task
:param TaskManager manager: The instance of
:py:class:`TaskManager` running this task.
Arguments passed to the constructor are available as `self.args`.
class StopTask(Task):
name = "stop_taskmanager"
def main(self, manager):
manager._running = False
class RateLimitContextManager:
def __init__(self, task_manager):
self.task_manager = task_manager
def __enter__(self):
if self.task_manager.last_ts is None:
while True:
delta = time.monotonic() - self.task_manager.last_ts
if delta >=
time.sleep( - delta)
def __exit__(self, etype, value, tb):
self.task_manager.last_ts = time.monotonic()
class TaskManager:
"""A single-threaded task dispatcher
This class is meant to be instantiated by a Provider in order to
execute remote API calls from a single thread with rate limiting.
:param str name: The name of the TaskManager (usually the provider name)
used in logging.
:param float rate_limit: The rate limit of the task manager expressed in
requests per second.
log = logging.getLogger("nodepool.driver.taskmanager.TaskManager")
def __init__(self, name, rate_limit):
self._running = True = name
self.queue = queue.Queue() = 1.0 / rate_limit
self.last_ts = None
def rateLimit(self):
"""Return a context manager to perform rate limiting. Use as follows:
.. code: python
with task_manager.rateLimit():
<execute API call>
return RateLimitContextManager(self)
def submitTask(self, task):
"""Submit a task to the task manager.
:param Task task: An instance of a subclass of :py:class:`Task`.
:returns: The submitted task for use in function chaning.
return task
def stop(self):
"""Stop the task manager."""
def run(self):
while True:
task = self.queue.get()
if not task:
self.log.debug("Manager %s running task %s (queue %s)" %
(,, self.queue.qsize()))
if not self._running:
except Exception:
self.log.exception("Task manager died")
class BaseTaskManagerProvider(Provider):
"""Subclass this to build a Provider with an included taskmanager"""
log = logging.getLogger("nodepool.driver.taskmanager.TaskManagerProvider")
def __init__(self, provider):
self.provider = provider
self.thread = None
self.task_manager = TaskManager(, provider.rate_limit)
def start(self, zk_conn):
if self.thread is None:
self.log.debug("Starting thread")
self.thread = threading.Thread(
def stop(self):
if self.thread is not None:
self.log.debug("Stopping thread")
def join(self):
if self.thread is not None: