from abc import ABCMeta, abstractmethod from collections import namedtuple import datetime import json import logging import re import time from enum import IntEnum, unique from six import add_metaclass, iteritems from trollius import async, coroutine, From, Return from urllib3.exceptions import ReadTimeoutError, ProtocolError import etcd import redis from buildman.asyncutil import wrap_with_threadpool from util import slash_join from util.expiresdict import ExpiresDict logger = logging.getLogger(__name__) ONE_DAY = 60 * 60 * 24 ORCHESTRATOR_UNAVAILABLE_SLEEP_DURATION = 5 DEFAULT_LOCK_EXPIRATION = 10000 ETCD_READ_TIMEOUT = 5 ETCD_MAX_WATCH_TIMEOUT = 30 REDIS_EXPIRING_SUFFIX = '/expiring' REDIS_DEFAULT_PUBSUB_KEY = 'orchestrator_events' REDIS_EVENT_KIND_MESSAGE = 'message' REDIS_EVENT_KIND_PMESSAGE = 'pmessage' REDIS_NONEXPIRING_KEY = -1 # This constant defines the Redis configuration flags used to watch [K]eyspace and e[x]pired # events on keys. For more info, see https://redis.io/topics/notifications#configuration REDIS_KEYSPACE_EVENT_CONFIG_VALUE = 'Kx' REDIS_KEYSPACE_EVENT_CONFIG_KEY = 'notify-keyspace-events' REDIS_KEYSPACE_KEY_PATTERN = '__keyspace@%s__:%s' REDIS_EXPIRED_KEYSPACE_PATTERN = slash_join(REDIS_KEYSPACE_KEY_PATTERN, REDIS_EXPIRING_SUFFIX) REDIS_EXPIRED_KEYSPACE_REGEX = re.compile(REDIS_EXPIRED_KEYSPACE_PATTERN % (r'(\S+)', r'(\S+)')) def orchestrator_from_config(manager_config, canceller_only=False): """ Allocates a new Orchestrator from the 'ORCHESTRATOR' block from provided manager config. Checks for legacy configuration prefixed with 'ETCD_' when the 'ORCHESTRATOR' is not present. :param manager_config: the configuration for the orchestrator :type manager_config: dict :rtype: :class: Orchestrator """ # Legacy codepath only knows how to configure etcd. if manager_config.get('ORCHESTRATOR') is None: manager_config['ORCHESTRATOR'] = {key: value for (key, value) in iteritems(manager_config) if key.startswith('ETCD_') and not key.endswith('_PREFIX')} # Sanity check that legacy prefixes are no longer being used. for key in manager_config['ORCHESTRATOR'].keys(): words = key.split('_') if len(words) > 1 and words[-1].lower() == 'prefix': raise AssertionError('legacy prefix used, use ORCHESTRATOR_PREFIX instead') def _dict_key_prefix(d): """ :param d: the dict that has keys prefixed with underscore :type d: {str: any} :rtype: str """ return d.keys()[0].split('_', 1)[0].lower() orchestrator_name = _dict_key_prefix(manager_config['ORCHESTRATOR']) def format_key(key): return key.lower().split('_', 1)[1] orchestrator_kwargs = {format_key(key): value for (key, value) in iteritems(manager_config['ORCHESTRATOR'])} if manager_config.get('ORCHESTRATOR_PREFIX') is not None: orchestrator_kwargs['orchestrator_prefix'] = manager_config['ORCHESTRATOR_PREFIX'] orchestrator_kwargs['canceller_only'] = canceller_only logger.debug('attempting to create orchestrator %s with kwargs %s', orchestrator_name, orchestrator_kwargs) return orchestrator_by_name(orchestrator_name, **orchestrator_kwargs) def orchestrator_by_name(name, **kwargs): _ORCHESTRATORS = { 'etcd': Etcd2Orchestrator, 'mem': MemoryOrchestrator, 'redis': RedisOrchestrator, } return _ORCHESTRATORS.get(name, MemoryOrchestrator)(**kwargs) class OrchestratorError(Exception): pass # TODO: replace with ConnectionError when this codebase is Python 3. class OrchestratorConnectionError(OrchestratorError): pass @unique class KeyEvent(IntEnum): CREATE = 1 SET = 2 DELETE = 3 EXPIRE = 4 class KeyChange(namedtuple('KeyChange', ['event', 'key', 'value'])): pass @add_metaclass(ABCMeta) class Orchestrator(object): """ Orchestrator is the interface that is used to synchronize the build states across build managers. This interface assumes that storage is being done by a key-value store that supports watching for events on keys. Missing keys should return KeyError; otherwise, errors should raise an OrchestratorError. :param key_prefix: the prefix of keys being watched :type key_prefix: str """ @abstractmethod def on_key_change(self, key, callback, restarter=None): """ The callback parameter takes in a KeyChange object as a parameter. """ pass @abstractmethod def get_prefixed_keys(self, prefix): """ :returns: a dict of key value pairs beginning with prefix :rtype: {str: str} """ pass @abstractmethod def get_key(self, key): """ :returns: the value stored at the provided key :rtype: str """ pass @abstractmethod def set_key(self, key, value, overwrite=False, expiration=None): """ :param key: the identifier for the value :type key: str :param value: the value being stored :type value: str :param overwrite: whether or not a KeyError is thrown if the key already exists :type overwrite: bool :param expiration: the duration in seconds that a key should be available :type expiration: int """ pass @abstractmethod def set_key_sync(self, key, value, overwrite=False, expiration=None): """ set_key, but without trollius coroutines. """ pass @abstractmethod def delete_key(self, key): """ Deletes a key that has been set in the orchestrator. :param key: the identifier for the key :type key: str """ pass @abstractmethod def lock(self, key, expiration=DEFAULT_LOCK_EXPIRATION): """ Takes a lock for synchronizing exclusive operations cluster-wide. :param key: the identifier for the lock :type key: str :param expiration: the duration until the lock expires :type expiration: :class:`datetime.timedelta` or int (seconds) :returns: whether or not the lock was acquired :rtype: bool """ pass @abstractmethod def shutdown(): """ This function should shutdown any final resources allocated by the Orchestrator. """ pass def _sleep_orchestrator(): """ This function blocks the trollius event loop by sleeping in order to backoff if a failure such as a ConnectionError has occurred. """ logger.exception('Connecting to etcd failed; sleeping for %s and then trying again', ORCHESTRATOR_UNAVAILABLE_SLEEP_DURATION) time.sleep(ORCHESTRATOR_UNAVAILABLE_SLEEP_DURATION) logger.exception('Connecting to etcd failed; slept for %s and now trying again', ORCHESTRATOR_UNAVAILABLE_SLEEP_DURATION) class EtcdAction(object): """ Enumeration of the various kinds of etcd actions we can observe via a watch. """ GET = 'get' SET = 'set' EXPIRE = 'expire' UPDATE = 'update' DELETE = 'delete' CREATE = 'create' COMPARE_AND_SWAP = 'compareAndSwap' COMPARE_AND_DELETE = 'compareAndDelete' class Etcd2Orchestrator(Orchestrator): def __init__(self, host='127.0.0.1', port=2379, cert_and_key=None, ca_cert=None, client_threads=5, canceller_only=False, **kwargs): self.is_canceller_only = canceller_only logger.debug('initializing async etcd client') self._sync_etcd_client = etcd.Client( host=host, port=port, cert=tuple(cert_and_key) if cert_and_key is not None else None, ca_cert=ca_cert, protocol='http' if cert_and_key is None else 'https', read_timeout=ETCD_READ_TIMEOUT, ) if not self.is_canceller_only: (self._etcd_client, self._async_executor) = wrap_with_threadpool(self._sync_etcd_client, client_threads) logger.debug('creating initial orchestrator state') self._shutting_down = False self._watch_tasks = {} @staticmethod def _sanity_check_ttl(ttl): """ A TTL of < 0 in etcd results in the key *never being expired*. We use a max here to ensure that if the TTL is < 0, the key will expire immediately. """ return max(ttl, 0) def _watch_etcd(self, key, callback, restarter=None, start_index=None): def callback_wrapper(changed_key_future): new_index = start_index etcd_result = None if not changed_key_future.cancelled(): try: etcd_result = changed_key_future.result() existing_index = getattr(etcd_result, 'etcd_index', None) new_index = etcd_result.modifiedIndex + 1 logger.debug('Got watch of key: %s at #%s with result: %s', key, existing_index, etcd_result) except ReadTimeoutError: logger.debug('Read-timeout on etcd watch %s, rescheduling', key) except etcd.EtcdEventIndexCleared: # This happens if etcd2 has moved forward too fast for us to start watching at the index # we retrieved. We therefore start a new watch at HEAD and (if specified) call the # restarter method which should conduct a read and reset the state of the manager. logger.debug('Etcd moved forward too quickly. Restarting watch cycle.') new_index = None if restarter is not None: async(restarter()) except (KeyError, etcd.EtcdKeyError): logger.debug('Etcd key already cleared: %s', key) return except etcd.EtcdConnectionFailed: _sleep_orchestrator() except etcd.EtcdException as eex: # TODO: This is a quick and dirty hack and should be replaced with a proper # exception check. if str(eex.message).find('Read timed out') >= 0: logger.debug('Read-timeout on etcd watch %s, rescheduling', key) else: logger.exception('Exception on etcd watch: %s', key) except ProtocolError: logger.exception('Exception on etcd watch: %s', key) if key not in self._watch_tasks or self._watch_tasks[key].done(): self._watch_etcd(key, callback, start_index=new_index, restarter=restarter) if etcd_result and etcd_result.value is not None: async(callback(self._etcd_result_to_keychange(etcd_result))) if not self._shutting_down: logger.debug('Scheduling watch of key: %s at start index %s', key, start_index) watch_future = self._etcd_client.watch(key, recursive=True, index=start_index, timeout=ETCD_MAX_WATCH_TIMEOUT) watch_future.add_done_callback(callback_wrapper) self._watch_tasks[key] = async(watch_future) @staticmethod def _etcd_result_to_keychange(etcd_result): event = Etcd2Orchestrator._etcd_result_to_keyevent(etcd_result) return KeyChange(event, etcd_result.key, etcd_result.value) @staticmethod def _etcd_result_to_keyevent(etcd_result): if etcd_result.action == EtcdAction.CREATE: return KeyEvent.CREATE if etcd_result.action == EtcdAction.SET: return KeyEvent.CREATE if etcd_result.createdIndex == etcd_result.modifiedIndex else KeyEvent.SET if etcd_result.action == EtcdAction.DELETE: return KeyEvent.DELETE if etcd_result.action == EtcdAction.EXPIRE: return KeyEvent.EXPIRE raise AssertionError('etcd action must have equivalant keyevent') def on_key_change(self, key, callback, restarter=None): assert not self.is_canceller_only logger.debug('creating watch on %s', key) self._watch_etcd(key, callback, restarter=restarter) @coroutine def get_prefixed_keys(self, prefix): assert not self.is_canceller_only try: etcd_result = yield From(self._etcd_client.read(prefix, recursive=True)) raise Return({leaf.key: leaf.value for leaf in etcd_result.leaves}) except etcd.EtcdKeyError: raise KeyError except etcd.EtcdConnectionFailed as ex: raise OrchestratorConnectionError(ex) except etcd.EtcdException as ex: raise OrchestratorError(ex) @coroutine def get_key(self, key): assert not self.is_canceller_only try: # Ignore pylint: the value property on EtcdResult is added dynamically using setattr. etcd_result = yield From(self._etcd_client.read(key)) raise Return(etcd_result.value) except etcd.EtcdKeyError: raise KeyError except etcd.EtcdConnectionFailed as ex: raise OrchestratorConnectionError(ex) except etcd.EtcdException as ex: raise OrchestratorError(ex) @coroutine def set_key(self, key, value, overwrite=False, expiration=None): assert not self.is_canceller_only yield From(self._etcd_client.write(key, value, prevExists=overwrite, ttl=self._sanity_check_ttl(expiration))) def set_key_sync(self, key, value, overwrite=False, expiration=None): self._sync_etcd_client.write(key, value, prevExists=overwrite, ttl=self._sanity_check_ttl(expiration)) @coroutine def delete_key(self, key): assert not self.is_canceller_only try: yield From(self._etcd_client.delete(key)) except etcd.EtcdKeyError: raise KeyError except etcd.EtcdConnectionFailed as ex: raise OrchestratorConnectionError(ex) except etcd.EtcdException as ex: raise OrchestratorError(ex) @coroutine def lock(self, key, expiration=DEFAULT_LOCK_EXPIRATION): assert not self.is_canceller_only try: yield From(self._etcd_client.write(key, {}, prevExist=False, ttl=self._sanity_check_ttl(expiration))) raise Return(True) except (KeyError, etcd.EtcdKeyError): raise Return(False) except etcd.EtcdConnectionFailed: logger.exception('Could not get etcd atomic lock as etcd is down') raise Return(False) except etcd.EtcdException as ex: raise OrchestratorError(ex) def shutdown(self): logger.debug('Shutting down etcd client.') self._shutting_down = True if self.is_canceller_only: return for (key, _), task in self._watch_tasks.items(): if not task.done(): logger.debug('Canceling watch task for %s', key) task.cancel() if self._async_executor is not None: self._async_executor.shutdown() class MemoryOrchestrator(Orchestrator): def __init__(self, **kwargs): self.state = ExpiresDict() self.callbacks = {} def _callbacks_prefixed(self, prefix): return (callback for (key, callback) in iteritems(self.callbacks) if key.startswith(prefix)) def on_key_change(self, key, callback, restarter=None): self.callbacks[key] = callback @coroutine def get_prefixed_keys(self, prefix): raise Return({k: value for (k, value) in self.state.items() if k.startswith(prefix)}) @coroutine def get_key(self, key): raise Return(self.state[key]) @coroutine def set_key(self, key, value, overwrite=False, expiration=None): preexisting_key = 'key' in self.state if preexisting_key and not overwrite: raise KeyError absolute_expiration = None if expiration is not None: absolute_expiration = datetime.datetime.now() + datetime.timedelta(seconds=expiration) self.state.set(key, value, expires=absolute_expiration) event = KeyEvent.CREATE if not preexisting_key else KeyEvent.SET for callback in self._callbacks_prefixed(key): yield From(callback(KeyChange(event, key, value))) def set_key_sync(self, key, value, overwrite=False, expiration=None): """ set_key, but without trollius coroutines. """ preexisting_key = 'key' in self.state if preexisting_key and not overwrite: raise KeyError absolute_expiration = None if expiration is not None: absolute_expiration = datetime.datetime.now() + datetime.timedelta(seconds=expiration) self.state.set(key, value, expires=absolute_expiration) event = KeyEvent.CREATE if not preexisting_key else KeyEvent.SET for callback in self._callbacks_prefixed(key): callback(KeyChange(event, key, value)) @coroutine def delete_key(self, key): value = self.state[key] del self.state[key] for callback in self._callbacks_prefixed(key): yield From(callback(KeyChange(KeyEvent.DELETE, key, value))) @coroutine def lock(self, key, expiration=DEFAULT_LOCK_EXPIRATION): if key in self.state: raise Return(False) self.state.set(key, None, expires=expiration) raise Return(True) def shutdown(self): self.state = None self.callbacks = None class RedisOrchestrator(Orchestrator): def __init__(self, host='127.0.0.1', port=6379, password=None, db=0, cert_and_key=None, ca_cert=None, client_threads=5, ssl=False, skip_keyspace_event_setup=False, canceller_only=False, **kwargs): self.is_canceller_only = canceller_only (cert, key) = tuple(cert_and_key) if cert_and_key is not None else (None, None) self._sync_client = redis.StrictRedis( host=host, port=port, password=password, db=db, ssl_certfile=cert, ssl_keyfile=key, ssl_ca_certs=ca_cert, ssl=ssl, ) self._shutting_down = False self._tasks = {} self._watched_keys = {} self._pubsub_key = slash_join(kwargs.get('orchestrator_prefix', ''), REDIS_DEFAULT_PUBSUB_KEY).lstrip('/') if not self.is_canceller_only: (self._client, self._async_executor) = wrap_with_threadpool(self._sync_client, client_threads) # Configure a subscription to watch events that the orchestrator manually publishes. logger.debug('creating pubsub with key %s', self._pubsub_key) published_pubsub = self._sync_client.pubsub() published_pubsub.subscribe(self._pubsub_key) (self._pubsub, self._async_executor_pub) = wrap_with_threadpool(published_pubsub) self._watch_published_key() # Configure a subscription to watch expired keyspace events. if not skip_keyspace_event_setup: self._sync_client.config_set(REDIS_KEYSPACE_EVENT_CONFIG_KEY, REDIS_KEYSPACE_EVENT_CONFIG_VALUE) expiring_pubsub = self._sync_client.pubsub() expiring_pubsub.psubscribe(REDIS_EXPIRED_KEYSPACE_PATTERN % (db, '*')) (self._pubsub_expiring, self._async_executor_ex) = wrap_with_threadpool(expiring_pubsub) self._watch_expiring_key() def _watch_published_key(self): def published_callback_wrapper(event_future): logger.debug('published callback called') event_result = None if not event_future.cancelled(): try: event_result = event_future.result() (redis_event, event_key, event_value) = event_result logger.debug('Got watch of key: (%s, %s, %s)', redis_event, event_key, event_value) except redis.ConnectionError: _sleep_orchestrator() except redis.RedisError: logger.exception('Exception watching redis publish: %s', event_key) # Schedule creating a new future if this one has been consumed. if 'pub' not in self._tasks or self._tasks['pub'].done(): self._watch_published_key() if event_result is not None and redis_event == REDIS_EVENT_KIND_MESSAGE: keychange = self._publish_to_keychange(event_value) for watched_key, callback in iteritems(self._watched_keys): if keychange.key.startswith(watched_key): async(callback(keychange)) if not self._shutting_down: logger.debug('Scheduling watch of publish stream') watch_future = self._pubsub.parse_response() watch_future.add_done_callback(published_callback_wrapper) self._tasks['pub'] = async(watch_future) def _watch_expiring_key(self): def expiring_callback_wrapper(event_future): logger.debug('expiring callback called') event_result = None if not event_future.cancelled(): try: event_result = event_future.result() if self._is_expired_keyspace_event(event_result): # Get the value of the original key before the expiration happened. key = self._key_from_expiration(event_future) expired_value = yield From(self._client.get(key)) # $KEY/expiring is gone, but the original key still remains, set an expiration for it # so that other managers have time to get the event and still read the expired value. yield From(self._client.expire(key, ONE_DAY)) except redis.ConnectionError: _sleep_orchestrator() except redis.RedisError: logger.exception('Exception watching redis expirations: %s', key) # Schedule creating a new future if this one has been consumed. if 'expire' not in self._tasks or self._tasks['expire'].done(): self._watch_expiring_key() if self._is_expired_keyspace_event(event_result) and expired_value is not None: for watched_key, callback in iteritems(self._watched_keys): if key.startswith(watched_key): async(callback(KeyChange(KeyEvent.EXPIRE, key, expired_value))) if not self._shutting_down: logger.debug('Scheduling watch of expiration') watch_future = self._pubsub_expiring.parse_response() watch_future.add_done_callback(expiring_callback_wrapper) self._tasks['expire'] = async(watch_future) def on_key_change(self, key, callback, restarter=None): assert not self.is_canceller_only logger.debug('watching key: %s', key) self._watched_keys[key] = callback @staticmethod def _is_expired_keyspace_event(event_result): """ Sanity check that this isn't an unrelated keyspace event. There could be a more efficient keyspace event config to avoid this client-side filter. """ if event_result is None: return False (redis_event, _pattern, matched_key, expired) = event_result return (redis_event == REDIS_EVENT_KIND_PMESSAGE and expired == 'expired' and REDIS_EXPIRED_KEYSPACE_REGEX.match(matched_key) is not None) @staticmethod def _key_from_expiration(event_result): (_redis_event, _pattern, matched_key, _expired) = event_result return REDIS_EXPIRED_KEYSPACE_REGEX.match(matched_key).groups()[1] @staticmethod def _publish_to_keychange(event_value): e = json.loads(event_value) return KeyChange(KeyEvent(e['event']), e['key'], e['value']) @coroutine def get_prefixed_keys(self, prefix): assert not self.is_canceller_only # TODO: This can probably be done with redis pipelines to make it transactional. keys = yield From(self._client.keys(prefix + '*')) # Yielding to the event loop is required, thus this cannot be written as a dict comprehension. results = {} for key in keys: if key.endswith(REDIS_EXPIRING_SUFFIX): continue ttl = yield From(self._client.ttl(key)) if ttl != REDIS_NONEXPIRING_KEY: # Only redis keys without expirations are live build manager keys. value = yield From(self._client.get(key)) results.update({key: value}) raise Return(results) @coroutine def get_key(self, key): assert not self.is_canceller_only value = yield From(self._client.get(key)) raise Return(value) @coroutine def set_key(self, key, value, overwrite=False, expiration=None): assert not self.is_canceller_only already_exists = yield From(self._client.exists(key)) yield From(self._client.set(key, value, xx=overwrite)) if expiration is not None: yield From(self._client.set(slash_join(key, REDIS_EXPIRING_SUFFIX), value, xx=overwrite, ex=expiration)) key_event = KeyEvent.SET if already_exists else KeyEvent.CREATE yield From(self._publish(event=key_event, key=key, value=value)) def set_key_sync(self, key, value, overwrite=False, expiration=None): already_exists = self._sync_client.exists(key) self._sync_client.set(key, value, xx=overwrite) if expiration is not None: self._sync_client.set(slash_join(key, REDIS_EXPIRING_SUFFIX), value, xx=overwrite, ex=expiration) self._sync_client.publish(self._pubsub_key, json.dumps({ 'event': int(KeyEvent.SET if already_exists else KeyEvent.CREATE), 'key': key, 'value': value, })) @coroutine def _publish(self, **kwargs): kwargs['event'] = int(kwargs['event']) event_json = json.dumps(kwargs) logger.debug('publishing event: %s', event_json) yield From(self._client.publish(self._pubsub_key, event_json)) @coroutine def delete_key(self, key): assert not self.is_canceller_only value = yield From(self._client.get(key)) yield From(self._client.delete(key)) yield From(self._client.delete(slash_join(key, REDIS_EXPIRING_SUFFIX))) yield From(self._publish(event=KeyEvent.DELETE, key=key, value=value)) @coroutine def lock(self, key, expiration=DEFAULT_LOCK_EXPIRATION): assert not self.is_canceller_only yield From(self.set_key(key, '', ex=expiration)) raise Return(True) @coroutine def shutdown(self): logger.debug('Shutting down redis client.') self._shutting_down = True if self.is_canceller_only: return for key, task in iteritems(self._tasks): if not task.done(): logger.debug('Canceling watch task for %s', key) task.cancel() if self._async_executor is not None: self._async_executor.shutdown() if self._async_executor_ex is not None: self._async_executor_ex.shutdown() if self._async_executor_pub is not None: self._async_executor_pub.shutdown()