class _ControllerProxy(ExportedState): def __init__(self, reactor, endpoint, elements, encoding, clock=None): self.__reactor = reactor self.__elements = elements self.__encoding = encoding self.__client_service = ClientService( endpoint=endpoint, factory=Factory.forProtocol(_ControllerProtocol), clock=clock) self.__client_service.startService() def state_def(self): for d in super(_ControllerProxy, self).state_def(): yield d for element in self.__elements: for d in IElement(element)._cells(self.__send, self.__encoding): yield d def close(self): """implements IComponent""" self.__client_service.stopService() def attach_context(self, device_context): """implements IComponent""" @defer.inlineCallbacks def __send(self, cmd): protocol = yield self.__client_service.whenConnected(failAfterFailures=1) # returned deferred is not actually used but in the future it would be good if Commands did in fact have a 'done' signal defer.returnValue(protocol.send(cmd))
class DeepstreamClient(Client): ''' This class instantiates an interface to interact with a Deepstream server. This class is the recommended mechanism for interacting with this module. It provides an interface to the other classes, each of which encapsulate a feature: Connection, Records, Events, RPC, and Presence. ''' def __init__(self, url=None, conn_string=None, authParams=None, reactor=None, **options): ''' Creates the client, but does not connect to the server automatically. Optional keyword parameters (**options) for... Client: url (required), authParams, reactor, conn_string, debug, factory protocol: url (required), authParams, heartbeat_interval rpc: rpcAckTimeout, rpcResponseTimeout, subscriptionTimeout record: recordReadAckTimeout, merge_strategy, recordReadTimeout, recordDeleteTimeout, recordDeepCopy, presence: subscriptionTimeout ''' if not url or url is None: raise ValueError( "url is None; you must specify a URL for the deepstream server, e.g. ws://localhost:6020/deepstream") parse_result = urlparse(url) if not authParams or authParams is None: authParams = {} if parse_result.username and parse_result.password: authParams['username'] = parse_result.username authParams['password'] = parse_result.password if not conn_string or conn_string is None: if parse_result.scheme == 'ws': if parse_result.hostname: conn_string = 'tcp:%s' % parse_result.hostname if parse_result.port: conn_string += ':%s' % parse_result.port else: conn_string += ':6020' if not conn_string or conn_string is None: raise ValueError( "Could not parse conn string from URL; you must specify a Twisted endpoint descriptor for the server, e.g. tcp:127.0.0.1:6020") if not reactor or reactor is None: from twisted.internet import reactor self.reactor = reactor factory = options.pop('factory', WSDeepstreamFactory) self._factory = factory(url, self, debug=options.pop('debug', False), reactor=reactor, **options) self._endpoint = clientFromString(reactor, conn_string) self._service = ClientService(self._endpoint, self._factory) # Handles reconnection for us EventEmitter.__init__(self) self._connection = ConnectionInterface(self, url) self._presence = PresenceHandler(self._connection, self, **options) self._event = EventHandler(self._connection, self, **options) self._rpc = RPCHandler(self._connection, self, **options) self._record = RecordHandler(self._connection, self, **options) self._message_callbacks = dict() self._message_callbacks[ constants.topic.PRESENCE] = self._presence.handle self._message_callbacks[ constants.topic.EVENT] = self._event.handle self._message_callbacks[ constants.topic.RPC] = self._rpc.handle self._message_callbacks[ constants.topic.RECORD] = self._record.handle self._message_callbacks[ constants.topic.ERROR] = self._on_error def login(self, auth_params): ''' Submit authentication credentials to the server once state is "Awaiting Authentication." Expects a dictionary. Options: User/pass: {'username': '******', 'password': '******'} Anonymous login/Open auth: {} https://deepstreamhub.com/tutorials/guides/open-auth/ Email: {'type': 'email', 'email': '*****@*****.**', 'password': '******'} https://deepstreamhub.com/tutorials/guides/email-auth/ Token: {'token': 'abcdefg'} https://deepstreamhub.com/tutorials/guides/token-auth/ Returns a Deferred ''' return self._connection.authenticate(auth_params) def connect(self, callback=None): ''' Connect to the server. Optionally, fire a callback once connected. Recommended callback is a login function. Calling the login function is automatic if the DeepstreamClient is instantiated with auth_params. ''' if callback: self._factory._connect_callback = callback if not self._service.running: self._service.startService() return def close(self): '''Legacy method: disconnect from the server.''' return self.disconnect() def disconnect(self): '''Terminate our connection to the server.''' # TODO: Say goodbye; clear message queue? self._factory._deliberate_close = True self._service.stopService() def whenAuthenticated(self, callback, *args): '''Execute a callback once authentication has succeeded.''' if self._factory._state == constants.connection_state.OPEN: callback(*args) else: self.once(constants.event.CONNECTION_STATE_CHANGED, lambda x: DeepstreamClient.whenAuthenticated(self, callback, *args)) # These properties are the same as in the parent class, but are repeated here for clarity @property def connection_state(self): return self._connection.state @property def record(self): return self._record @property def event(self): return self._event @property def rpc(self): return self._rpc @property def presence(self): return self._presence
def stopService(self): try: yield ClientService.stopService(self) except Exception as e: log.failure("Exception {excp!s}", excp=e) reactor.stop()
class RemoteHandler(logging.Handler): """ An logging handler that writes logs to a TCP target. Args: host: The host of the logging target. port: The logging target's port. maximum_buffer: The maximum buffer size. """ def __init__( self, host: str, port: int, maximum_buffer: int = 1000, level=logging.NOTSET, _reactor=None, ): super().__init__(level=level) self.host = host self.port = port self.maximum_buffer = maximum_buffer self._buffer: Deque[logging.LogRecord] = deque() self._connection_waiter: Optional[Deferred] = None self._producer: Optional[LogProducer] = None # Connect without DNS lookups if it's a direct IP. if _reactor is None: from twisted.internet import reactor _reactor = reactor try: ip = ip_address(self.host) if isinstance(ip, IPv4Address): endpoint: IStreamClientEndpoint = TCP4ClientEndpoint( _reactor, self.host, self.port ) elif isinstance(ip, IPv6Address): endpoint = TCP6ClientEndpoint(_reactor, self.host, self.port) else: raise ValueError("Unknown IP address provided: %s" % (self.host,)) except ValueError: endpoint = HostnameEndpoint(_reactor, self.host, self.port) factory = Factory.forProtocol(Protocol) self._service = ClientService(endpoint, factory, clock=_reactor) self._service.startService() self._stopping = False self._connect() def close(self): self._stopping = True self._service.stopService() def _connect(self) -> None: """ Triggers an attempt to connect then write to the remote if not already writing. """ # Do not attempt to open multiple connections. if self._connection_waiter: return def fail(failure: Failure) -> None: # If the Deferred was cancelled (e.g. during shutdown) do not try to # reconnect (this will cause an infinite loop of errors). if failure.check(CancelledError) and self._stopping: return # For a different error, print the traceback and re-connect. failure.printTraceback(file=sys.__stderr__) self._connection_waiter = None self._connect() def writer(result: Protocol) -> None: # Force recognising transport as a Connection and not the more # generic ITransport. transport: Connection = result.transport # type: ignore # We have a connection. If we already have a producer, and its # transport is the same, just trigger a resumeProducing. if self._producer and transport is self._producer.transport: self._producer.resumeProducing() self._connection_waiter = None return # If the producer is still producing, stop it. if self._producer: self._producer.stopProducing() # Make a new producer and start it. self._producer = LogProducer( buffer=self._buffer, transport=transport, format=self.format, ) transport.registerProducer(self._producer, True) self._producer.resumeProducing() self._connection_waiter = None deferred: Deferred = self._service.whenConnected(failAfterFailures=1) deferred.addCallbacks(writer, fail) self._connection_waiter = deferred def _handle_pressure(self) -> None: """ Handle backpressure by shedding records. The buffer will, in this order, until the buffer is below the maximum: - Shed DEBUG records. - Shed INFO records. - Shed the middle 50% of the records. """ if len(self._buffer) <= self.maximum_buffer: return # Strip out DEBUGs self._buffer = deque( filter(lambda record: record.levelno > logging.DEBUG, self._buffer) ) if len(self._buffer) <= self.maximum_buffer: return # Strip out INFOs self._buffer = deque( filter(lambda record: record.levelno > logging.INFO, self._buffer) ) if len(self._buffer) <= self.maximum_buffer: return # Cut the middle entries out buffer_split = floor(self.maximum_buffer / 2) old_buffer = self._buffer self._buffer = deque() for _ in range(buffer_split): self._buffer.append(old_buffer.popleft()) end_buffer = [] for _ in range(buffer_split): end_buffer.append(old_buffer.pop()) self._buffer.extend(reversed(end_buffer)) def emit(self, record: logging.LogRecord) -> None: self._buffer.append(record) # Handle backpressure, if it exists. try: self._handle_pressure() except Exception: # If handling backpressure fails, clear the buffer and log the # exception. self._buffer.clear() logger.warning("Failed clearing backpressure") # Try and write immediately. self._connect()
class TerseJSONToTCPLogObserver(object): """ An IObserver that writes JSON logs to a TCP target. Args: hs (HomeServer): The homeserver that is being logged for. host: The host of the logging target. port: The logging target's port. metadata: Metadata to be added to each log entry. """ hs = attr.ib() host = attr.ib(type=str) port = attr.ib(type=int) metadata = attr.ib(type=dict) maximum_buffer = attr.ib(type=int) _buffer = attr.ib(default=attr.Factory(deque), type=deque) _connection_waiter = attr.ib(default=None, type=Optional[Deferred]) _logger = attr.ib(default=attr.Factory(Logger)) _producer = attr.ib(default=None, type=Optional[LogProducer]) def start(self) -> None: # Connect without DNS lookups if it's a direct IP. try: ip = ip_address(self.host) if isinstance(ip, IPv4Address): endpoint = TCP4ClientEndpoint(self.hs.get_reactor(), self.host, self.port) elif isinstance(ip, IPv6Address): endpoint = TCP6ClientEndpoint(self.hs.get_reactor(), self.host, self.port) except ValueError: endpoint = HostnameEndpoint(self.hs.get_reactor(), self.host, self.port) factory = Factory.forProtocol(Protocol) self._service = ClientService(endpoint, factory, clock=self.hs.get_reactor()) self._service.startService() self._connect() def stop(self): self._service.stopService() def _connect(self) -> None: """ Triggers an attempt to connect then write to the remote if not already writing. """ if self._connection_waiter: return self._connection_waiter = self._service.whenConnected( failAfterFailures=1) @self._connection_waiter.addErrback def fail(r): r.printTraceback(file=sys.__stderr__) self._connection_waiter = None self._connect() @self._connection_waiter.addCallback def writer(r): # We have a connection. If we already have a producer, and its # transport is the same, just trigger a resumeProducing. if self._producer and r.transport is self._producer.transport: self._producer.resumeProducing() self._connection_waiter = None return # If the producer is still producing, stop it. if self._producer: self._producer.stopProducing() # Make a new producer and start it. self._producer = LogProducer(buffer=self._buffer, transport=r.transport) r.transport.registerProducer(self._producer, True) self._producer.resumeProducing() self._connection_waiter = None def _handle_pressure(self) -> None: """ Handle backpressure by shedding events. The buffer will, in this order, until the buffer is below the maximum: - Shed DEBUG events - Shed INFO events - Shed the middle 50% of the events. """ if len(self._buffer) <= self.maximum_buffer: return # Strip out DEBUGs self._buffer = deque( filter(lambda event: event["level"] != "DEBUG", self._buffer)) if len(self._buffer) <= self.maximum_buffer: return # Strip out INFOs self._buffer = deque( filter(lambda event: event["level"] != "INFO", self._buffer)) if len(self._buffer) <= self.maximum_buffer: return # Cut the middle entries out buffer_split = floor(self.maximum_buffer / 2) old_buffer = self._buffer self._buffer = deque() for i in range(buffer_split): self._buffer.append(old_buffer.popleft()) end_buffer = [] for i in range(buffer_split): end_buffer.append(old_buffer.pop()) self._buffer.extend(reversed(end_buffer)) def __call__(self, event: dict) -> None: flattened = flatten_event(event, self.metadata, include_time=True) self._buffer.append(flattened) # Handle backpressure, if it exists. try: self._handle_pressure() except Exception: # If handling backpressure fails,clear the buffer and log the # exception. self._buffer.clear() self._logger.failure("Failed clearing backpressure") # Try and write immediately. self._connect()
class BaseNetwork: def __init__(self, ledger): self.config = ledger.config self.client = None self.service = None self.running = False self._on_connected_controller = StreamController() self.on_connected = self._on_connected_controller.stream self._on_header_controller = StreamController() self.on_header = self._on_header_controller.stream self._on_status_controller = StreamController() self.on_status = self._on_status_controller.stream self.subscription_controllers = { 'blockchain.headers.subscribe': self._on_header_controller, 'blockchain.address.subscribe': self._on_status_controller, } @defer.inlineCallbacks def start(self): for server in cycle(self.config['default_servers']): connection_string = 'tcp:{}:{}'.format(*server) endpoint = clientFromString(reactor, connection_string) log.debug("Attempting connection to SPV wallet server: %s", connection_string) self.service = ClientService(endpoint, StratumClientFactory(self)) self.service.startService() try: self.client = yield self.service.whenConnected(failAfterFailures=2) yield self.ensure_server_version() log.info("Successfully connected to SPV wallet server: %s", connection_string) self._on_connected_controller.add(True) yield self.client.on_disconnected.first except CancelledError: return except Exception: # pylint: disable=broad-except log.exception("Connecting to %s raised an exception:", connection_string) finally: self.client = None if not self.running: return def stop(self): self.running = False if self.service is not None: self.service.stopService() if self.is_connected: return self.client.on_disconnected.first else: return defer.succeed(True) @property def is_connected(self): return self.client is not None and self.client.connected def rpc(self, list_or_method, *args): if self.is_connected: return self.client.rpc(list_or_method, *args) else: raise ConnectionError("Attempting to send rpc request when connection is not available.") def ensure_server_version(self, required='1.2'): return self.rpc('server.version', __version__, required) def broadcast(self, raw_transaction): return self.rpc('blockchain.transaction.broadcast', raw_transaction) def get_history(self, address): return self.rpc('blockchain.address.get_history', address) def get_transaction(self, tx_hash): return self.rpc('blockchain.transaction.get', tx_hash) def get_merkle(self, tx_hash, height): return self.rpc('blockchain.transaction.get_merkle', tx_hash, height) def get_headers(self, height, count=10000): return self.rpc('blockchain.block.headers', height, count) def subscribe_headers(self): return self.rpc('blockchain.headers.subscribe', True) def subscribe_address(self, address): return self.rpc('blockchain.address.subscribe', address)
class ZenHubClient(object): """A client for connecting to ZenHub as a ZenHub Worker. After start is called, this class automatically handles connecting to ZenHub, registering the zenhubworker with ZenHub, and automatically reconnecting to ZenHub if the connection to ZenHub is corrupted for any reason. """ def __init__( self, reactor, endpoint, credentials, worker, timeout, worklistId, ): """Initialize a ZenHubClient instance. :type reactor: IReactorCore :param endpoint: Where zenhub is found :type endpoint: IStreamClientEndpoint :param credentials: Credentials to log into ZenHub. :type credentials: IUsernamePassword :param worker: Reference to worker :type worker: IReferenceable :param float timeout: Seconds to wait before determining whether ZenHub is unresponsive. :param str worklistId: Name of the worklist to receive tasks from. """ self.__reactor = reactor self.__endpoint = endpoint self.__credentials = credentials self.__worker = worker self.__timeout = timeout self.__worklistId = worklistId self.__stopping = False self.__pinger = None self.__service = None self.__log = getLogger(self) self.__signalFile = ConnectedToZenHubSignalFile() def start(self): """Start connecting to ZenHub.""" self.__stopping = False factory = ZenPBClientFactory() self.__service = ClientService( self.__endpoint, factory, retryPolicy=backoffPolicy(initialDelay=0.5, factor=3.0), ) self.__service.startService() self.__prepForConnection() def stop(self): """Stop connecting to ZenHub.""" self.__stopping = True self.__reset() def restart(self): """Restart the connect to ZenHub.""" self.__reset() self.start() def __reset(self): if self.__pinger: self.__pinger.stop() self.__pinger = None if self.__service: self.__service.stopService() self.__service = None self.__signalFile.remove() def __prepForConnection(self): if not self.__stopping: self.__log.info("Prepping for connection") self.__service.whenConnected().addCallbacks( self.__connected, self.__notConnected, ) def __disconnected(self, *args): # Called when the connection to ZenHub is lost. # Ensures that processing resumes when the connection to ZenHub # is restored. self.__log.info( "Lost connection to ZenHub: %s", args[0] if args else "<no reason given>", ) if self.__pinger: self.__pinger.stop() self.__pinger = None self.__signalFile.remove() self.__prepForConnection() def __notConnected(self, *args): self.__log.info("Not connected! %r", args) @defer.inlineCallbacks def __connected(self, broker): # Called when a connection to ZenHub is established. # Logs into ZenHub and passes up a worker reference for ZenHub # to use to dispatch method calls. # Sometimes broker.transport doesn't have a 'socket' attribute if not hasattr(broker.transport, "socket"): self.restart() defer.returnValue(None) self.__log.info("Connection to ZenHub established") try: setKeepAlive(broker.transport.socket) zenhub = yield self.__login(broker) yield zenhub.callRemote( "reportingForWork", self.__worker, workerId=self.__worker.instanceId, worklistId=self.__worklistId, ) ping = PingZenHub(zenhub, self) self.__pinger = task.LoopingCall(ping) d = self.__pinger.start(self.__timeout, now=False) d.addErrback(self.__pingFail) # Catch and pass on errors except defer.CancelledError: self.__log.error("Timed out trying to login to ZenHub") self.restart() defer.returnValue(None) except Exception as ex: self.__log.error( "Unable to report for work: (%s) %s", type(ex).__name__, ex, ) self.__signalFile.remove() self.__reactor.stop() else: self.__log.info("Logged into ZenHub") self.__signalFile.touch() # Connection complete; install a listener to be notified if # the connection is lost. broker.notifyOnDisconnect(self.__disconnected) def __login(self, broker): d = broker.factory.login(self.__credentials, self.__worker) timeoutCall = self.__reactor.callLater(self.__timeout, d.cancel) def completedLogin(arg): if timeoutCall.active(): timeoutCall.cancel() return arg d.addBoth(completedLogin) return d def __pingFail(self, ex): self.__log.error("Pinger failed: %s", ex)
class NotificationSourceIntegrationTest(IntegrationTest): @inlineCallbacks def setUp(self): super(NotificationSourceIntegrationTest, self).setUp() self.endpoint = AMQEndpoint( reactor, self.rabbit.config.hostname, self.rabbit.config.port, username="******", password="******", heartbeat=1) self.policy = backoffPolicy(initialDelay=0) self.factory = AMQFactory(spec=AMQP0_8_SPEC_PATH) self.service = ClientService( self.endpoint, self.factory, retryPolicy=self.policy) self.connector = NotificationConnector(self.service) self.source = NotificationSource(self.connector) self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") self.service.startService() @inlineCallbacks def tearDown(self): self.service.stopService() super(NotificationSourceIntegrationTest, self).tearDown() # Wrap resetting queues and client in a try/except, since the broker # may have been stopped (e.g. when this is the last test being run). try: yield self.channel.queue_delete(queue="uuid") except: pass finally: yield self.client.close() @inlineCallbacks def test_get_after_publish(self): """ Calling get() after a message has been published in the associated queue returns a Notification for that message. """ yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_before_publish(self): """ Calling get() before a message has been published in the associated queue will wait until publication. """ deferred = self.source.get("uuid", 0) self.assertFalse(deferred.called) yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield deferred self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_with_error(self): """ If an error occurs in during get(), the client is closed so we can query messages again. """ yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) with self.assertRaises(NotFound): yield self.source.get("uuid-unknown", 0) notification = yield self.source.get("uuid", 0) self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_concurrent_with_error(self): """ If an error occurs in a call to get(), other calls don't fail, and are retried on reconnection instead. """ client1 = yield self.service.whenConnected() deferred = self.source.get("uuid", 0) with self.assertRaises(NotFound): yield self.source.get("uuid-unknown", 0) yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield deferred self.assertEqual("hello", notification.payload) client2 = yield self.service.whenConnected() # The ClientService has reconnected, yielding a new client. self.assertIsNot(client1, client2) @inlineCallbacks def test_get_timeout(self): """ Calls to get() timeout after a certain amount of time if no message arrived on the queue. """ self.source.timeout = 1 with self.assertRaises(Timeout): yield self.source.get("uuid", 0) client = yield self.service.whenConnected() channel = yield client.channel(1) # The channel is still opened self.assertFalse(channel.closed) # The consumer has been deleted self.assertNotIn("uuid.0", client.queues) @inlineCallbacks def test_get_with_broker_shutdown_during_consume(self): """ If rabbitmq gets shutdown during the basic-consume call, we wait for the reconection and retry transparently. """ # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. yield self.connector() d = self.source.get("uuid", 0) # Restart rabbitmq yield self.client.close() yield self.client.disconnected.wait() self.rabbit.cleanUp() self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp() # Get a new channel and re-declare the queue, since the restart has # destroyed it. self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") # Publish a message in the queue yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield d self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_with_broker_die_during_consume(self): """ If rabbitmq dies during the basic-consume call, we wait for the reconection and retry transparently. """ # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. yield self.connector() d = self.source.get("uuid", 0) # Kill rabbitmq and start it again yield self.client.close() yield self.client.disconnected.wait() self.rabbit.runner.kill() self.rabbit.cleanUp() self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp() # Get a new channel and re-declare the queue, since the crash has # destroyed it. self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") # Publish a message in the queue yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield d self.assertEqual("hello", notification.payload) @inlineCallbacks def test_wb_get_with_broker_shutdown_during_message_wait(self): """ If rabbitmq gets shutdown while we wait for messages, we transparently wait for the reconnection and try again. """ # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. yield self.connector() d = self.source.get("uuid", 0) # Acquiring the channel lock makes sure that basic-consume has # succeeded and we started waiting for the message. yield self.source._channel_lock.acquire() self.source._channel_lock.release() # Restart rabbitmq yield self.client.close() yield self.client.disconnected.wait() self.rabbit.cleanUp() self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp() # Get a new channel and re-declare the queue, since the restart has # destroyed it. self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") # Publish a message in the queue yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield d self.assertEqual("hello", notification.payload) @inlineCallbacks def test_wb_heartbeat(self): """ If heartbeat checks fail due to network issues, we keep re-trying until the network recovers. """ self.service.stopService() # Put a TCP proxy between NotificationSource and RabbitMQ, to simulate # packets getting dropped on the floor. proxy = ProxyService( self.rabbit.config.hostname, self.rabbit.config.port) proxy.startService() self.addCleanup(proxy.stopService) self.endpoint._port = proxy.port self.service = ClientService( self.endpoint, self.factory, retryPolicy=self.policy) self.connector._service = self.service self.service.startService() # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. channel = yield self.connector() deferred = self.source.get("uuid", 0) # Start dropping packets on the floor proxy.block() # Publish a notification, which won't be delivered just yet. yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) # Wait for the first connection to terminate, because heartbeat # checks will fail. yield channel.client.disconnected.wait() # Now let packets flow again. proxy.unblock() # The situation got recovered. notification = yield deferred self.assertEqual("hello", notification.payload) self.assertEqual(2, proxy.connections) @inlineCallbacks def test_reject_notification(self): """ Calling reject() on a Notification puts the associated message back in the queue so that it's available to subsequent get() calls. """ yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) yield notification.reject() notification = yield self.source.get("uuid", 1) self.assertEqual("hello", notification.payload) @inlineCallbacks def test_ack_message(self): """ Calling ack() on a Notification confirms the removal of the associated message from the queue, making subsequent calls waiting for another message. """ yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) yield notification.ack() yield self.channel.basic_publish( routing_key="uuid", content=Content("hello 2")) notification = yield self.source.get("uuid", 1) self.assertEqual("hello 2", notification.payload) @inlineCallbacks def test_ack_with_broker_shutdown(self): """ If rabbitmq gets shutdown before we ack a Notification, an error is raised. """ client = yield self.service.whenConnected() yield self.channel.basic_publish( routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) self.rabbit.cleanUp() yield client.disconnected.wait() try: yield notification.ack() except Bounced: pass else: self.fail("Notification not bounced") self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp()
class ZenHubClient(object): """A client for connecting to ZenHub as a ZenHub Worker. After start is called, this class automatically handles connecting to ZenHub, registering the zenhubworker with ZenHub, and automatically reconnecting to ZenHub if the connection to ZenHub is corrupted for any reason. """ def __init__(self, reactor, endpoint, credentials, worker, timeout): """Initialize a ZenHubClient instance. @param reactor {IReactorCore} @param endpoint {IStreamClientEndpoint} Where zenhub is found @param credentials {IUsernamePassword} Credentials to log into ZenHub. @param worker {IReferenceable} Reference to worker @param timeout {float} Seconds to wait before determining whether ZenHub is unresponsive. """ self.__reactor = reactor self.__endpoint = endpoint self.__credentials = credentials self.__worker = worker self.__timeout = timeout self.__stopping = False self.__pinger = None self.__service = None self.__log = getLogger(self) self.__signalFile = ConnectedToZenHubSignalFile() def start(self): """Start connecting to ZenHub.""" self.__stopping = False factory = ZenPBClientFactory() self.__service = ClientService( self.__endpoint, factory, retryPolicy=backoffPolicy(initialDelay=0.5, factor=3.0), ) self.__service.startService() self.__prepForConnection() def stop(self): """Stop connecting to ZenHub.""" self.__stopping = True self.__reset() def restart(self): """Restart the connect to ZenHub.""" self.__reset() self.start() def __reset(self): if self.__pinger: self.__pinger.stop() self.__pinger = None if self.__service: self.__service.stopService() self.__service = None self.__signalFile.remove() def __prepForConnection(self): if not self.__stopping: self.__log.info("Prepping for connection") self.__service.whenConnected().addCallbacks( self.__connected, self.__notConnected, ) def __disconnected(self, *args): # Called when the connection to ZenHub is lost. # Ensures that processing resumes when the connection to ZenHub # is restored. self.__log.info( "Lost connection to ZenHub: %s", args[0] if args else "<no reason given>", ) if self.__pinger: self.__pinger.stop() self.__pinger = None self.__signalFile.remove() self.__prepForConnection() def __notConnected(self, *args): self.__log.info("Not connected! %r", args) @defer.inlineCallbacks def __connected(self, broker): # Called when a connection to ZenHub is established. # Logs into ZenHub and passes up a worker reference for ZenHub # to use to dispatch method calls. # Sometimes broker.transport doesn't have a 'socket' attribute if not hasattr(broker.transport, "socket"): self.restart() defer.returnValue(None) self.__log.info("Connection to ZenHub established") try: setKeepAlive(broker.transport.socket) zenhub = yield self.__login(broker) yield zenhub.callRemote( "reportingForWork", self.__worker, workerId=self.__worker.instanceId, ) ping = PingZenHub(zenhub, self) self.__pinger = task.LoopingCall(ping) d = self.__pinger.start(self.__timeout, now=False) d.addErrback(self.__pingFail) # Catch and pass on errors except defer.CancelledError: self.__log.error("Timed out trying to login to ZenHub") self.restart() defer.returnValue(None) except Exception as ex: self.__log.error( "Unable to report for work: (%s) %s", type(ex), ex, ) self.__signalFile.remove() self.__reactor.stop() else: self.__log.info("Logged into ZenHub") self.__signalFile.touch() # Connection complete; install a listener to be notified if # the connection is lost. broker.notifyOnDisconnect(self.__disconnected) def __login(self, broker): d = broker.factory.login(self.__credentials, self.__worker) timeoutCall = self.__reactor.callLater(self.__timeout, d.cancel) def completedLogin(arg): if timeoutCall.active(): timeoutCall.cancel() return arg d.addBoth(completedLogin) return d def __pingFail(self, ex): self.__log.error("Pinger failed: %s", ex)
class MQTTService(object): """MQTT Service interface to Azure IoT hub. Attributes: client: (ClientService): Twisted client service connected (bool): Service connection flag devid (str): Device identifer username: (str): Azure IoT Hub MQTT username password: (str): Azure IoT Hub MQTT password messages (list): Received inbound messages """ TIMEOUT = 10.0 def __init__(self, endpoint, factory, devid, username, password): self.client = ClientService(endpoint, factory) self.connected = False self.devid = devid self.username = username self.password = password self.messages = [] @inlineCallbacks def publishMessage(self, data): """Publish the MQTT message. Any inbound messages are copied to the messages list attribute, and returned to the caller. Args: data (str): Application data to send Returns: A list of received messages. """ # Start the service, and add a timeout to check the connection. self.client.startService() reactor.callLater(self.TIMEOUT, self.checkConnection) # Attempt to connect. If we tiemout and cancel and exception # is thrown. try: yield self.client.whenConnected().addCallback( self.azureConnect, data) except Exception as e: log.error("Azure MQTT service failed to connect to broker.") # Stop the service if sucessful, and finally return # any inbound messages. else: yield self.client.stopService() finally: returnValue(self.messages) @inlineCallbacks def checkConnection(self): """Check if the connected flag is set. Stop the service if not. """ if not self.connected: yield self.client.stopService() @inlineCallbacks def azureConnect(self, protocol, data): self.connected = True protocol.setWindowSize(1) protocol.onPublish = self.onPublish pubtopic = 'devices/{}/messages/events/'.format(self.devid) subtopic = 'devices/{}/messages/devicebound/#'.format(self.devid) try: # Connect and subscribe yield protocol.connect(self.devid, username=self.username, password=self.password, cleanStart=False, keepalive=10) yield protocol.subscribe(subtopic, 2) except Exception as e: log.error( "Azure MQTT service could not connect to " "Azure IOT Hub using username {name}", name=self.username) returnValue(None) # Publish the outbound message yield protocol.publish(topic=pubtopic, qos=0, message=str(data)) def onPublish(self, topic, payload, qos, dup, retain, msgId): """Receive messages from Azure IoT Hub IoT Hub delivers messages with the Topic Name devices/{device_id}/messages/devicebound/ or devices/{device_id}/messages/devicebound/{property_bag} if there are any message properties. {property_bag} contains url-encoded key/value pairs of message properties. System property names have the prefix $, application properties use the original property name with no prefix. """ message = '' # Split the component parameters of topic. Obtain the downstream message # using the key name message. params = parse_qs(topic) if 'message' in params: self.messages.append(params['message'])
class NotificationSourceIntegrationTest(IntegrationTest): @inlineCallbacks def setUp(self): super(NotificationSourceIntegrationTest, self).setUp() self.endpoint = AMQEndpoint(reactor, self.rabbit.config.hostname, self.rabbit.config.port, username="******", password="******", heartbeat=1) self.policy = backoffPolicy(initialDelay=0) self.factory = AMQFactory(spec=AMQP0_8_SPEC_PATH) self.service = ClientService(self.endpoint, self.factory, retryPolicy=self.policy) self.connector = NotificationConnector(self.service) self.source = NotificationSource(self.connector) self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") self.service.startService() @inlineCallbacks def tearDown(self): self.service.stopService() super(NotificationSourceIntegrationTest, self).tearDown() # Wrap resetting queues and client in a try/except, since the broker # may have been stopped (e.g. when this is the last test being run). try: yield self.channel.queue_delete(queue="uuid") except: pass finally: yield self.client.close() @inlineCallbacks def test_get_after_publish(self): """ Calling get() after a message has been published in the associated queue returns a Notification for that message. """ yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_before_publish(self): """ Calling get() before a message has been published in the associated queue will wait until publication. """ deferred = self.source.get("uuid", 0) self.assertFalse(deferred.called) yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield deferred self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_with_error(self): """ If an error occurs in during get(), the client is closed so we can query messages again. """ yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) with self.assertRaises(NotFound): yield self.source.get("uuid-unknown", 0) notification = yield self.source.get("uuid", 0) self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_concurrent_with_error(self): """ If an error occurs in a call to get(), other calls don't fail, and are retried on reconnection instead. """ client1 = yield self.service.whenConnected() deferred = self.source.get("uuid", 0) with self.assertRaises(NotFound): yield self.source.get("uuid-unknown", 0) yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield deferred self.assertEqual("hello", notification.payload) client2 = yield self.service.whenConnected() # The ClientService has reconnected, yielding a new client. self.assertIsNot(client1, client2) @inlineCallbacks def test_get_timeout(self): """ Calls to get() timeout after a certain amount of time if no message arrived on the queue. """ self.source.timeout = 1 with self.assertRaises(Timeout): yield self.source.get("uuid", 0) client = yield self.service.whenConnected() channel = yield client.channel(1) # The channel is still opened self.assertFalse(channel.closed) # The consumer has been deleted self.assertNotIn("uuid.0", client.queues) @inlineCallbacks def test_get_with_broker_shutdown_during_consume(self): """ If rabbitmq gets shutdown during the basic-consume call, we wait for the reconection and retry transparently. """ # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. yield self.connector() d = self.source.get("uuid", 0) # Restart rabbitmq yield self.client.close() yield self.client.disconnected.wait() self.rabbit.cleanUp() self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp() # Get a new channel and re-declare the queue, since the restart has # destroyed it. self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") # Publish a message in the queue yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield d self.assertEqual("hello", notification.payload) @inlineCallbacks def test_get_with_broker_die_during_consume(self): """ If rabbitmq dies during the basic-consume call, we wait for the reconection and retry transparently. """ # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. yield self.connector() d = self.source.get("uuid", 0) # Kill rabbitmq and start it again yield self.client.close() yield self.client.disconnected.wait() self.rabbit.runner.kill() self.rabbit.cleanUp() self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp() # Get a new channel and re-declare the queue, since the crash has # destroyed it. self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") # Publish a message in the queue yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield d self.assertEqual("hello", notification.payload) @inlineCallbacks def test_wb_get_with_broker_shutdown_during_message_wait(self): """ If rabbitmq gets shutdown while we wait for messages, we transparently wait for the reconnection and try again. """ # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. yield self.connector() d = self.source.get("uuid", 0) # Acquiring the channel lock makes sure that basic-consume has # succeeded and we started waiting for the message. yield self.source._channel_lock.acquire() self.source._channel_lock.release() # Restart rabbitmq yield self.client.close() yield self.client.disconnected.wait() self.rabbit.cleanUp() self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp() # Get a new channel and re-declare the queue, since the restart has # destroyed it. self.client = yield self.endpoint.connect(self.factory) self.channel = yield self.client.channel(1) yield self.channel.channel_open() yield self.channel.queue_declare(queue="uuid") # Publish a message in the queue yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield d self.assertEqual("hello", notification.payload) @inlineCallbacks def test_wb_heartbeat(self): """ If heartbeat checks fail due to network issues, we keep re-trying until the network recovers. """ self.service.stopService() # Put a TCP proxy between NotificationSource and RabbitMQ, to simulate # packets getting dropped on the floor. proxy = ProxyService(self.rabbit.config.hostname, self.rabbit.config.port) proxy.startService() self.addCleanup(proxy.stopService) self.endpoint._port = proxy.port self.service = ClientService(self.endpoint, self.factory, retryPolicy=self.policy) self.connector._service = self.service self.service.startService() # This will make the connector setup the channel before we call # get(), so by the time we call it in the next line all # connector-related deferreds will fire synchronously and the # code will block on basic-consume. channel = yield self.connector() deferred = self.source.get("uuid", 0) # Start dropping packets on the floor proxy.block() # Publish a notification, which won't be delivered just yet. yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) # Wait for the first connection to terminate, because heartbeat # checks will fail. yield channel.client.disconnected.wait() # Now let packets flow again. proxy.unblock() # The situation got recovered. notification = yield deferred self.assertEqual("hello", notification.payload) self.assertEqual(2, proxy.connections) @inlineCallbacks def test_reject_notification(self): """ Calling reject() on a Notification puts the associated message back in the queue so that it's available to subsequent get() calls. """ yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) yield notification.reject() notification = yield self.source.get("uuid", 1) self.assertEqual("hello", notification.payload) @inlineCallbacks def test_ack_message(self): """ Calling ack() on a Notification confirms the removal of the associated message from the queue, making subsequent calls waiting for another message. """ yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) yield notification.ack() yield self.channel.basic_publish(routing_key="uuid", content=Content("hello 2")) notification = yield self.source.get("uuid", 1) self.assertEqual("hello 2", notification.payload) @inlineCallbacks def test_ack_with_broker_shutdown(self): """ If rabbitmq gets shutdown before we ack a Notification, an error is raised. """ client = yield self.service.whenConnected() yield self.channel.basic_publish(routing_key="uuid", content=Content("hello")) notification = yield self.source.get("uuid", 0) self.rabbit.cleanUp() yield client.disconnected.wait() try: yield notification.ack() except Bounced: pass else: self.fail("Notification not bounced") self.rabbit.config = RabbitServerResources( port=self.rabbit.config.port) # Ensure that we use the same port self.rabbit.setUp()
def stopService(self): ClientService.stopService(self)