class Zmq_broker(BaseModule): context = None s_pub = None pub_endpoint = None serialize_to = None serialize = None def __init__(self, mod_conf, pub_endpoint, serialize_to): from zmq import Context, PUB BaseModule.__init__(self, mod_conf) self.pub_endpoint = pub_endpoint self.serialize_to = serialize_to logger.info("[Zmq Broker] Binding to endpoint " + self.pub_endpoint) # This doesn't work properly in init() # sometimes it ends up beings called several # times and the address becomes already in use. self.context = Context() self.s_pub = self.context.socket(PUB) self.s_pub.bind(self.pub_endpoint) # Load the correct serialization function # depending on the serialization method # chosen in the configuration. if self.serialize_to == "msgpack": from msgpack import Packer packer = Packer(default=encode_monitoring_data) self.serialize = lambda msg: packer.pack(msg) elif self.serialize_to == "json": self.serialize = lambda msg: json.dumps(msg, cls=SetEncoder) else: raise Exception("[Zmq Broker] No valid serialization method defined (Got " + str(self.serialize_to) + ")!") # Called by Broker to say 'let's prepare yourself guy' def init(self): logger.info("[Zmq Broker] Initialization of the Zmq broker module") # Publish to the ZeroMQ socket # using the chosen serialization method def publish(self, msg, topic=""): from zmq import SNDMORE data = self.serialize(msg) self.s_pub.send(topic, SNDMORE) self.s_pub.send(data) # An host check have just arrived, we UPDATE data info with this def manage_brok(self, b): logger.debug("[Zmq Broker] Got broker update: " + str(b.data)) # Publish update data to the ZeroMQ endpoint. msg = b.data self.publish(msg, b.type) # Properly close down this thing. def do_stop(self): self.s_pub.close() self.context.term()
class Listener(Thread): def __init__(self): super(Listener, self).__init__(name="Listener") self._shutdown = False self.context = Context() self.sub = self.context.socket(SUB) self.sub.bind('tcp://*:7000') self.sub.setsockopt(SUBSCRIBE, "") self.poller = Poller() self.poller.register(self.sub, POLLIN) def cleanup(self): self.sub.close() self.context.term() def run(self): while True: socks = dict(self.poller.poll(timeout=1)) if socks.get(self.sub) == POLLIN: msg = self.sub.recv(flags=NOBLOCK) print msg if self._shutdown: break self.cleanup()
class HomeBase(Thread): def __init__(self): super(HomeBase, self).__init__(name="HomeBase") self.context = Context() self.pull = self.context.socket(PULL) self.pull.bind("tcp://*:7001") self._shutdown = False self.poller = Poller() self.poller.register(self.pull, POLLIN) def cleanup(self): self.pull.close() self.context.term() def run(self): while True: socks = dict(self.poller.poll(timeout=1)) if socks.get(self.pull) == POLLIN: msg = self.pull.recv() msg += ", WORK RECEIVED " print msg if self._shutdown: break self.cleanup()
class Leatherneck(Thread): def __init__(self): super(Leatherneck, self).__init__(name="Leatherneck") self.context = Context() self.pull = self.context.socket(PULL) self.pull.connect("tcp://localhost:7000") self.push = self.context.socket(PUSH) self.push.connect("tcp://localhost:7001") self.poller = Poller() self.poller.register(self.pull, POLLIN) self._shutdown = False for th in t_enum(): if th.name == "MainThread": self.mainthread = th def cleanup(self): print "Workers exiting..." self.push.close() self.pull.close() self.context.term() def run(self): while True: if not self.mainthread.is_alive(): self._shutdown = True break socks = dict(self.poller.poll(timeout=1)) if socks.get(self.pull) == POLLIN: msg = self.pull.recv(flags=NOBLOCK) msg += " WORK COMPLETE, " + str(time()) self.push.send(msg, flags=NOBLOCK) if self._shutdown: break self.cleanup()
class Leatherneck(Thread): def __init__(self): super(Leatherneck, self).__init__(name="Leatherneck") self.context = Context() self.pull = self.context.socket(PULL) self.pull.connect("tcp://localhost:7000") self.push = self.context.socket(PUSH) self.push.connect("tcp://localhost:7001") self.poller = Poller() self.poller.register(self.pull, POLLIN) self._shutdown = False def cleanup(self): self.push.close() self.pull.close() self.context.term() def run(self): while True: socks = dict(self.poller.poll(timeout=1)) if socks.get(self.pull) == POLLIN: msg = self.pull.recv() msg += " WORK COMPLETE, " + str(time()) self.push.send(msg) if self._shutdown: break self.cleanup()
class DrillingWell(Thread): def __init__(self): super(DrillingWell, self).__init__(name="DrillingWell") self.context = Context() self.push = self.context.socket(PUSH) self.push.bind("tcp://*:7000") self._shutdown = False for th in t_enum(): if th.name == "MainThread": self.mainthread = th def cleanup(self): print "Producer exiting..." self.push.close() self.context.term() def run(self): count = 0 while True: if not self.mainthread.is_alive(): self._shutdown = True break sleep(0.01) count += 1 self.push.send("SOMETHING " + str(count)) if self._shutdown: break self.cleanup()
class HomeBase(Thread): def __init__(self): super(HomeBase, self).__init__(name="HomeBase") self.context = Context() self.pull = self.context.socket(PULL) self.pull.bind("tcp://*:7001") self._shutdown = False self.poller = Poller() self.poller.register(self.pull, POLLIN) for th in t_enum(): if th.name == "MainThread": self.mainthread = th def cleanup(self): print "Home exiting..." self.pull.close() self.context.term() def run(self): while True: if not self.mainthread.is_alive(): self._shutdown = True break socks = dict(self.poller.poll(timeout=1)) if socks.get(self.pull) == POLLIN: msg = self.pull.recv(flags=NOBLOCK) msg += ", WORK RECEIVED " print msg if self._shutdown: break self.cleanup()
class Zmq_broker(BaseModule): context = None s_pub = None pub_endpoint = None serialize_to = None serialize = None def __init__(self, mod_conf, pub_endpoint, serialize_to): from zmq import Context, PUB BaseModule.__init__(self, mod_conf) self.pub_endpoint = pub_endpoint self.serialize_to = serialize_to logger.info("[Zmq Broker] Binding to endpoint " + self.pub_endpoint) # This doesn't work properly in init() # sometimes it ends up beings called several # times and the address becomes already in use. self.context = Context() self.s_pub = self.context.socket(PUB) self.s_pub.bind(self.pub_endpoint) # Load the correct serialization function # depending on the serialization method # chosen in the configuration. if self.serialize_to == "msgpack": from msgpack import Packer packer = Packer(default=encode_monitoring_data) self.serialize = lambda msg: packer.pack(msg) elif self.serialize_to == "json": self.serialize = lambda msg: json.dumps(msg, cls=SetEncoder) else: raise Exception( "[Zmq Broker] No valid serialization method defined (Got " + str(self.serialize_to) + ")!") # Called by Broker to say 'let's prepare yourself guy' def init(self): logger.info("[Zmq Broker] Initialization of the Zmq broker module") # Publish to the ZeroMQ socket # using the chosen serialization method def publish(self, msg, topic=""): from zmq import SNDMORE data = self.serialize(msg) self.s_pub.send(topic, SNDMORE) self.s_pub.send(data) # An host check have just arrived, we UPDATE data info with this def manage_brok(self, b): logger.debug("[Zmq Broker] Got broker update: " + str(b.data)) #Publish update data to the ZeroMQ endpoint. msg = b.data self.publish(msg, b.type) # Properly close down this thing. def do_stop(self): self.s_pub.close() self.context.term()
class ZmqFactory(object): """ I control individual ZeroMQ connections. Factory creates and destroys ZeroMQ context. :var reactor: reference to Twisted reactor used by all the connections :var ioThreads: number of IO threads ZeroMQ will be using for this context :vartype ioThreads: int :var lingerPeriod: number of milliseconds to block when closing socket (terminating context), when there are some messages pending to be sent :vartype lingerPeriod: int :var connections: set of instanciated :class:`ZmqConnection` :vartype connections: set :var context: ZeroMQ context """ reactor = reactor ioThreads = 1 lingerPeriod = 100 def __init__(self): """ Constructor. Create ZeroMQ context. """ self.connections = set() self.context = Context(self.ioThreads) def __repr__(self): return "ZmqFactory()" def shutdown(self): """ Shutdown factory. This is shutting down all created connections and terminating ZeroMQ context. Also cleans up Twisted reactor. """ for connection in self.connections.copy(): connection.shutdown() self.connections = None self.context.term() self.context = None def registerForShutdown(self): """ Register factory to be automatically shut down on reactor shutdown. It is recommended that this method is called on any created factory. """ reactor.addSystemEventTrigger('during', 'shutdown', self.shutdown)
class ZmqSubscriber(HiddenSubscriber): """ Subscriber class subscribing to a certain topic Attributes: context (zmq.Context): socket (Socket): Socket object of ZMQ context topic (String): Topic subscriber subscribes to """ def __init__(self, url, topic): """ Initializes object Args: url (String): url to publish messages to topic (String): Topic to publish messages under """ super(ZmqSubscriber, self).__init__(url) self._context = Context() self._socket = self._context.socket(SUB) self._socket.setsockopt(SUBSCRIBE, topic) self._socket.setsockopt(RCVTIMEO, 500) # Wait 500ms for message to arrive self._socket.connect(url) self._logger = logging.getLogger('ZeromqSubscriber') def receive(self): """ Receives a message Returns: String """ topic, message = self._socket.recv_multipart() return message def __enter__(self): """ Statement used for the `` with ... as ...:`` returns the object to use in the ``with`` block Returns: ZmqSubscriber """ return self def __exit__(self, exc_type, exc_value, exc_tb): """ Executed when leaving ``with`` block, regardless whether because of an exception or normal program flow """ self._socket.close() self._context.term()
class ZmqPublisher(HiddenPublisher): """ Publisher class publishing messages to a certain topic to an url Attributes: context (zmq.Context): socket (Socket): Socket object of ZMQ context topic (String): Topic publisher publishs to """ def __init__(self, url, topic): """ Initializes object Args: url (String): url to publish messages to topic (String): Topic to publish messages under """ super(ZmqPublisher, self).__init__(topic) self._context = Context() self._socket = self._context.socket(PUB) self._socket.bind(url) self._logger = logging.getLogger('ZeromqPublisher') def publish(self, message): """ Publishes message Args: message (String): Message to publish """ self._socket.send_multipart([self.topic, message]) def __enter__(self): """ Statement used for the `` with ... as ...:`` returns the object to use in the ``with`` block Returns: ZmqPublisher """ return self def __exit__(self, exc_type, exc_value, exc_tb): """ Executed when leaving ``with`` block, regardless whether because of an exception or normal program flow """ self._socket.close() self._context.term()
class Server: """Server component handling zmq requests.""" def __init__( self, host: str, port: str, calls: Dict[str, Tuple[Callable[[Body], Response], Optional[Dict]]], io_threads: int = 1, ) -> None: """Initialize a Server with a host, port and calls.""" self._calls = calls self._host = host self._port = port self._init_server(io_threads) def _init_server(self, io_threads: int) -> None: self._context = Context(io_threads=io_threads) self._socket = self._context.socket(REP) self._socket.bind("tcp://{:s}:{:s}".format(self._host, self._port)) def start(self) -> None: """Start the server loop.""" while True: request: Request = self._socket.recv_json() response: Response = self._handle_request(request) self._socket.send_json(response) def _handle_request(self, request: Request) -> Response: try: validate(instance=request, schema=request_schema) func, schema = self._calls[request["header"]["message"]] if schema: validate(request["body"], schema=schema) return func(request["body"]) except ValidationError: return get_response(400) except KeyError: return get_response(404) def close(self) -> None: """Close the socket and terminate it.""" self._socket.close() self._context.term()
class ZmqContext(object): """Provides a singleton wrapper for a ZeroMQ context""" self_ = None def __init__(self, iothreads): assert not ZmqContext.self_ self._context = Context(iothreads) def __del__(self): self._context.term() def socket(self, socket_type): """Creates and returns a socket of the given type""" return self._context.socket(socket_type) @staticmethod def instance(iothreads=4): """Returns the singleton instance of the ZeroMQ context""" if not ZmqContext.self_: ZmqContext.self_ = ZmqContext(iothreads) return ZmqContext.self_
print "%s exiting..." % self.name finally: self.cleanup() context = Context() heart = Heartbeat(context, name="Heartbeat Thread") stethoscope = Stethoscope(context, name="Stethoscope Thread") for t in (heart, stethoscope): t.start() while True: try: # call thread.join to keep some control in the main thread while (heart.is_alive() or stethoscope.is_alive()): heart.join(timeout=0.1) stethoscope.join(timeout=0.1) except KeyboardInterrupt: shutdown.set() while (heart.is_alive() or stethoscope.is_alive()): heart.join(timeout=0.1) stethoscope.join(timeout=0.1) context.term() break
def close_zmq_ctx(ctx: zmq.Context): ctx.destroy() ctx.term()
class GuldAI(object): # TODO bug: does not know which .git/ tree is active (use git.toplevel) # TODO bug: does not read .gitignore (use fnmatch from stdlib) # TODO push to all remotes # TODO feature: post-process fsync's (GPG, AES, torrent) # TODO feature: pre-process on open/read (GPG, AES, torrent) # TODO assume always create git submodules # TODO feature: publish to TPC port # TODO feature: listen for, validate, and fetch valid commits from friends # TODO feature: listen for, validate and act on commands from ADMIN stop = None def __init__(self, name=None): self.name = name logging.basicConfig(filename='/var/log/guldfs/gai-%s.log' % self.name, level=logging.DEBUG) self.context = Context() self.subscriber = self.context.socket(SUB) self.subscriber.connect("ipc:///tmp/guldfs0.ipc") self.subscriber.setsockopt(SUBSCRIBE, b"guldfs") # publisher = context.socket(zmq.PUB) # publisher.bind("tpc://127.0.0.1:6070") self.stop = False self.repos = {} def run(self): # TODO replace with more efficient loop while not self.stop: # Read envelope with address [channel, message] = self.subscriber.recv_multipart() self.handle_fs_message(channel, message) logging.info("done running") def handle_fs_message(self, channel, message): if (channel != b'guldfs'): return parts = message.decode().split(":") # TODO be more specific about this if hasattr(self, parts[0]) and callable(getattr(self, parts[0])): args = [] if len(parts) > 1: args = parts[1:] getattr(self, parts[0])(*args) #else: # logging.warn('unknown message: %s' % message) def get_or_load_repo(self, path): if path in self.repos: return self.repos[path] else: self.repos[path] = Git(path) return self.repos[path] def mount(self, mountpoint, user): repo = self.get_or_load_repo('/') repo.pull() def destroy(self): logging.info('destroying') # clean up zmq socket connections self.subscriber.close() self.context.term() self.stop = True def fsync(self, path): if '.git/' not in path and not path.startswith('/git/'): repo = self.get_or_load_repo(os.path.dirname(path)) repo.add(cfg.rawpath(path).replace(repo.toplevel, '')) repo.commit('gai: found changes to %s' % path) def fsyncdir(self, path): # TODO add files, commit and push. Cascade up to user home. if '.git/' not in path and not path.startswith('/git/'): repo = self.get_or_load_repo(path) repo.stash() repo.pull() repo.stash('apply') #repo.commit('gai: found changes to %s' % path) def open(self, path, flags): if '.git/' not in path and not path.startswith('/git/'): # TODO pull! check signatures! decrypt and otherwise post-process repo = get_or_load_repo(path) repo.pull() def readdir(self, path): if '.git/' not in path and not path.startswith('/git/'): # TODO pull! check signatures! decrypt and otherwise post-process repo = get_or_load_repo(path) repo.pull()
class TestClient(unittest.TestCase): def setUp(self): self.context = Context() self.start_router_sockets() def tearDown(self): self.close_sockets() self.context.term() def start_router_sockets(self): self.api = self.context.socket(REP) self.broadcast = self.context.socket(PUB) self.api.bind(API_BIND_ADDRESS) self.broadcast.bind(BROADCAST_BIND_ADDRESS) def close_sockets(self): self.api.close() self.broadcast.close() def test_connect_raises_ValueError_when_no_communication_channel_is_specified(self): client = Client() with self.assertRaises(ValueError): client.connect() def test_api_methods_should_raise_RuntimeError_if_not_connected_to_api(self): client = Client() client.connect(broadcast=BROADCAST_ADDRESS) with self.assertRaises(RuntimeError): client.send_api_request({'command': 'get configuration'}) with self.assertRaises(RuntimeError): client.get_api_reply() with self.assertRaises(RuntimeError): client.api_poll(timeout=1) # milliseconds with self.assertRaises(RuntimeError): client.disconnect_api() def test_broadcast_methods_should_raise_RuntimeError_if_not_connected_to_broadcast(self): client = Client() client.connect(api=API_ADDRESS) with self.assertRaises(RuntimeError): client.broadcast_subscribe('42') with self.assertRaises(RuntimeError): client.broadcast_unsubscribe('42') with self.assertRaises(RuntimeError): client.broadcast_poll(timeout=1) # milliseconds with self.assertRaises(RuntimeError): client.broadcast_receive() with self.assertRaises(RuntimeError): client.disconnect_broadcast() def test_send_api_request(self): client = Client() client.connect(api=API_ADDRESS) client.send_api_request({'command': 'get configuration'}) if not self.api.poll(TIMEOUT): self.fail('Timeout wainting for API command') message = self.api.recv_json() self.assertEqual(message, {'command': 'get configuration'}) def test_get_api_reply(self): client = Client() client.connect(api=API_ADDRESS) client.send_api_request({'command': 'get configuration'}) if not self.api.poll(TIMEOUT): self.fail('Timeout wainting for API command') self.api.recv_json() self.api.send_json({'configuration': 'spam eggs ham'}) message = client.get_api_reply() # what if it hangs? self.assertEqual(message, {'configuration': 'spam eggs ham'}) def test_api_poll(self, timeout=0): client = Client() client.connect(api=API_ADDRESS) client.send_api_request({'command': 'get configuration'}) if not self.api.poll(TIMEOUT): self.fail('Timeout waiting for API message') self.api.recv_json() start_time = time.time() result = client.api_poll(TIMEOUT) end_time = time.time() self.assertFalse(result) # there is no message, should wait for the entire TIMEOUT total_time = (end_time - start_time) * 1000 # milliseconds self.assertTrue(TIMEOUT <= total_time <= 1.1 * TIMEOUT) self.api.send_json({'configuration': 'spam eggs ham'}) start_time = time.time() result = client.api_poll(TIMEOUT) end_time = time.time() self.assertTrue(result) # poll should return almost immediatly (there is a message) total_time = (end_time - start_time) * 1000 # milliseconds self.assertTrue(total_time < TIMEOUT) def test_broadcast_subscribe_poll_and_receive(self): client = Client() client.connect(broadcast=BROADCAST_ADDRESS) client.broadcast_subscribe('spam') time.sleep(TIMEOUT / 1000.0) # wait for subscribe to take effect self.broadcast.send('spam eggs ham') start_time = time.time() poll_result = client.broadcast_poll(TIMEOUT) end_time = time.time() self.assertTrue(poll_result) total_time = (end_time - start_time) * 1000 self.assertTrue(total_time < TIMEOUT) message = client.broadcast_receive() # what if it hangs? self.assertEqual(message, 'spam eggs ham') self.broadcast.send('eggs ham') start_time = time.time() poll_result = client.broadcast_poll(TIMEOUT) end_time = time.time() self.assertFalse(poll_result) total_time = (end_time - start_time) * 1000 self.assertTrue(TIMEOUT <= total_time <= 1.1 * TIMEOUT) def test_broadcast_unsubscribe(self): client = Client() client.connect(broadcast=BROADCAST_ADDRESS) client.broadcast_subscribe('spam') time.sleep(TIMEOUT / 1000.0) # wait for subscribe to take effect self.broadcast.send('spam eggs ham') self.assertTrue(client.broadcast_poll(TIMEOUT)) self.assertEqual(client.broadcast_receive(), 'spam eggs ham') client.broadcast_unsubscribe('spam') self.broadcast.send('spam eggs ham') self.assertFalse(client.broadcast_poll(TIMEOUT)) def test_disconnect(self): client = Client() client.connect(api=API_ADDRESS, broadcast=BROADCAST_ADDRESS) #connected we can communicate... client.send_api_request({'command': 'get configuration'}) self.assertTrue(self.api.poll(TIMEOUT)) self.api.recv_json() self.api.send_json({'command': 'ok'}) self.assertTrue(client.api_poll(TIMEOUT)) self.assertEqual(client.get_api_reply(), {'command': 'ok'}) client.broadcast_subscribe('spam') time.sleep(TIMEOUT / 1000.0) # wait for subscribe to take effect self.broadcast.send('spam eggs ham') self.assertTrue(client.broadcast_poll(TIMEOUT)) self.assertEqual(client.broadcast_receive(), 'spam eggs ham') #disconnected not! client.disconnect_api() with self.assertRaises(RuntimeError): client.send_api_request({'command': 'get configuration'}) with self.assertRaises(RuntimeError): client.api_poll(TIMEOUT) with self.assertRaises(RuntimeError): client.get_api_reply() client.broadcast_subscribe('spam') time.sleep(TIMEOUT / 1000.0) # wait for subscribe to take effect self.broadcast.send('spam eggs ham') self.assertTrue(client.broadcast_poll(TIMEOUT)) self.assertEqual(client.broadcast_receive(), 'spam eggs ham') client.disconnect_broadcast() with self.assertRaises(RuntimeError): client.broadcast_subscribe('spam') with self.assertRaises(RuntimeError): client.broadcast_poll(TIMEOUT) with self.assertRaises(RuntimeError): client.broadcast_receive() #connect again... client.connect(api=API_ADDRESS, broadcast=BROADCAST_ADDRESS) client.send_api_request({'command': 'get configuration'}) self.assertTrue(self.api.poll(TIMEOUT)) self.api.recv_json() self.api.send_json({'command': 'ok'}) self.assertTrue(client.api_poll(TIMEOUT)) self.assertEqual(client.get_api_reply(), {'command': 'ok'}) client.broadcast_subscribe('spam') time.sleep(TIMEOUT / 1000.0) # wait for subscribe to take effect self.broadcast.send('spam eggs ham') self.assertTrue(client.broadcast_poll(TIMEOUT)) self.assertEqual(client.broadcast_receive(), 'spam eggs ham') #disconnected everything client.disconnect() with self.assertRaises(RuntimeError): client.send_api_request({'command': 'get configuration'}) with self.assertRaises(RuntimeError): client.api_poll(TIMEOUT) with self.assertRaises(RuntimeError): client.get_api_reply() with self.assertRaises(RuntimeError): client.broadcast_subscribe('spam') with self.assertRaises(RuntimeError): client.broadcast_poll(TIMEOUT) with self.assertRaises(RuntimeError): client.broadcast_receive() with self.assertRaises(RuntimeError): client.disconnect() # Should not raises: client.disconnect(silent=True) client.disconnect_api(silent=True) client.disconnect_broadcast(silent=True)
class WorkloadGenerator(object): """Object responsible for generating workload.""" def __init__( self, generator_listening: str, generator_port: str, workload_listening: str, workload_pub_port: str, ) -> None: """Initialize a WorkloadGenerator.""" self._workload_listening = workload_listening self._workload_pub_port = workload_pub_port server_calls: Dict = { "start workload": self._call_start_workload, "get workload": self._call_get_workload, "stop workload": self._call_stop_workload, } self._server = Server(generator_listening, generator_port, server_calls) self._workload: str = None # type: ignore self._workload_frequency: int = 0 self._init_server() self._init_scheduler() def _init_scheduler(self) -> None: self._scheduler = BackgroundScheduler() self._generate_workload_job = self._scheduler.add_job( func=self._generate_workload, trigger="interval", seconds=1, ) self._scheduler.start() def __enter__(self) -> "WorkloadGenerator": """Return self for a context manager.""" return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: """Call close with a context manager.""" self.close() return None def _init_server(self) -> None: self._context = Context(io_threads=1) self._pub_socket = self._context.socket(PUB) self._pub_socket.bind("tcp://{:s}:{:s}".format( self._workload_listening, self._workload_pub_port)) def _call_start_workload(self, body: Body) -> Response: self._workload = body["workload_name"] self._workload_frequency = body["frequency"] return get_response(200) def _call_stop_workload(self, body: Body) -> Response: self._workload = None # type: ignore self._workload_frequency = 0 return get_response(200) def _generate_workload(self) -> None: if self._workload: response = get_response(200) response["body"]["querylist"] = [ self._workload for _ in range(self._workload_frequency) ] self._pub_socket.send_json(response) def _call_get_workload(self, body: Body) -> Response: response = get_response(200) response["body"]["workload"] = { "workload_name": self._workload, "frequency": self._workload_frequency, } return response def start(self) -> None: """Start the generator by starting the server.""" self._server.start() def close(self) -> None: """Close the socket and context.""" self._generate_workload_job.remove() self._scheduler.shutdown() self._pub_socket.close() self._context.term()
class ZmqContextManager(object): """ I control individual ZeroMQ connections. Factory creates and destroys ZeroMQ context. :var reactor: reference to Twisted reactor used by all the connections :var ioThreads: number of IO threads ZeroMQ will be using for this context :vartype ioThreads: int :var lingerPeriod: number of milliseconds to block when closing socket (terminating context), when there are some messages pending to be sent :vartype lingerPeriod: int :var connections: set of instantiated :class:`ZmqConnection` :vartype connections: set :var context: ZeroMQ context """ reactor = reactor ioThreads = 1 lingerPeriod = 100 _instance = None def __new__(cls): if not cls._instance: cls._instance = super(ZmqContextManager, cls).__new__(cls) cls._instance.initialized = False return cls._instance def __init__(self): """ Constructor. Create ZeroMQ context. """ if not self.initialized: self.initialized = True self.connections = set() self.context = Context(self.ioThreads) reactor.addSystemEventTrigger('during', 'shutdown', self.shutdown) def __repr__(self): return "ZmqContextManager(%d threads)" % self.ioThreads def shutdown(self): """ Shutdown factory. This is shutting down all created connections and terminating ZeroMQ context. Also cleans up Twisted reactor. """ if not self.initialized: return self.initialized = False for connection in self.connections.copy(): connection.shutdown() self.connections = None self.context.term() self.context = None
raise KillThread except KillThread: print "%s exiting..." % self.name finally: self.cleanup() context = Context() heart = Heartbeat(context, name="Heartbeat Thread") stethoscope = Stethoscope(context, name="Stethoscope Thread") for t in (heart, stethoscope): t.start() while True: try: # call thread.join to keep some control in the main thread while (heart.is_alive() or stethoscope.is_alive()): heart.join(timeout=0.1) stethoscope.join(timeout=0.1) except KeyboardInterrupt: shutdown.set() while (heart.is_alive() or stethoscope.is_alive()): heart.join(timeout=0.1) stethoscope.join(timeout=0.1) context.term() break
#!/usr/bin/env python from json import dumps from msgpack import packb, unpackb from sys import argv from zmq import Context, REQ ctx = Context() sock = ctx.socket(REQ) sock.connect("ipc:///var/run/pcma.socket") sock.send(packb(argv[1:])) ret, out = unpackb(sock.recv(), encoding='utf-8') if ret: print(dumps(out)) else: raise Exception(out) sock.close() ctx.term()
class WorkloadGenerator(object): """Object responsible for generating workload.""" def __init__( self, generator_listening: str, generator_port: str, workload_listening: str, workload_pub_port: str, default_workload_location: str, ) -> None: """Initialize a WorkloadGenerator.""" self._workload_listening = workload_listening self._workload_pub_port = workload_pub_port self._default_workload_location = default_workload_location server_calls: Dict[str, Tuple[Callable[[Body], Response], Optional[Dict]]] = { "start workload": ( self._call_start_workload, start_workload_request_schema, ), "stop workload": (self._call_stop_workload, None), } self._server = Server(generator_listening, generator_port, server_calls) self._generate_workload_flag = False self._frequency = 0 self._workloads: Dict[str, Any] = {} self._init_server() self._init_scheduler() def _init_scheduler(self) -> None: self._scheduler = BackgroundScheduler() self._generate_workload_job = self._scheduler.add_job( func=self._generate_workload, trigger="interval", seconds=1, ) self._scheduler.start() def __enter__(self) -> "WorkloadGenerator": """Return self for a context manager.""" return self def __exit__( self, exc_type: Optional[Type[BaseException]], exc_value: Optional[BaseException], traceback: Optional[TracebackType], ) -> Optional[bool]: """Call close with a context manager.""" self.close() return None def _init_server(self) -> None: self._context = Context(io_threads=1) self._pub_socket = self._context.socket(PUB) self._pub_socket.bind("tcp://{:s}:{:s}".format( self._workload_listening, self._workload_pub_port)) def _get_default_workload_location(self) -> str: return self._default_workload_location def _get_workload(self, workload_type: str) -> Workload: workload = self._workloads.get(workload_type) if not workload: workload = Workload(workload_type, self._get_default_workload_location()) self._workloads[workload_type] = workload return workload def _call_start_workload(self, body: Body) -> Response: frequency: int = body["frequency"] workload_type: str = body["folder_name"] try: self._get_workload(workload_type) except ( NotExistingWorkloadFolderException, EmptyWorkloadFolderException, QueryTypeNotFoundException, QueryTypesNotSpecifiedException, ) as e: return get_error_response(400, str(e)) self._workload_type = workload_type self._frequency = frequency self._generate_workload_flag = True return get_response(200) def _call_stop_workload(self, body: Body) -> Response: self._generate_workload_flag = False return get_response(200) def _publish_data(self, data: Response) -> None: self._pub_socket.send_json(data) def _generate_workload(self) -> None: if self._generate_workload_flag: workload = self._get_workload(self._workload_type) queries = workload.generate_workload(self._frequency) response = get_response(200) response["body"] = {"querylist": queries} self._publish_data(response) def start(self) -> None: """Start the generator by starting the server.""" self._server.start() def close(self) -> None: """Close the socket and context.""" self._generate_workload_job.remove() self._scheduler.shutdown() self._pub_socket.close() self._context.term()
class TranslationsClient(): """A translation service client. Not thread safe! Create and use one instance per thread. """ def __init__(self, server_host, server_port, timeout=3, encoding="utf-8", skip_translations=False, translation_cache=None): """ Create a client instance which can be used to get translations. A client will close all resources on exit. If resource should be freed before that use the close function. :type server_host: str :type server_port: int :param timeout: Timeout for a request, in seconds. :type timeout: int :type encoding: str """ self._timeout = timeout or 0.001 # seconds self._encoding = encoding self._context = Context(1) self._socket = self._context.socket(DEALER) self._socket.set(LINGER, 0) # do not wait for anything when closing self._socket.connect("tcp://{}:{}".format(server_host, server_port)) self._poller = Poller() self._poller.register(self._socket, POLLIN) self._skip_translations = skip_translations self._translation_cache = translation_cache if self._translation_cache and not isinstance(self._translation_cache, TranslationCache): raise TypeError( "The cache needs to be a subclass of TranslationCache") def _handle_response(self, req_id, response): """ Check the response and extract the translations. :type response: [bytes] :rtype: [str] or None """ translations = None if len(response) < 3: raise TranslationsServerError("Server could not handle request.") else: response = [r.decode(self._encoding) for r in response] response_id, _ = response[:2] if response_id != req_id: _LOG.info("Got a response for an old or unknown request.", extra={"response": response}) else: translations = response[2:] if translations == [""]: raise TranslationsServerError( "Server encountered an error.") return translations def get(self, lang, country, *keys): """ Try to find translations for the given keys (and plural forms). If the service is not available or a translation can not be found the keys are returned instead. Note: if only one key was requested, only a string will be returned instead of a list of strings! :param keys: Keys and or tuples of key and plural to translate. :type keys: str or (str, int) :return List of translations for each key or key and plural tuple. If something can not be translated (or the service is down) the keys will be returned instead of a translation. If there was only one key or tuple requested than only that translation is returned instead of a list of translations. :rtype: str or [str] """ if not keys: raise ValueError("No keys to translate!") if self._context is None: raise AttributeError("Client is closed.") keys = [(k, "") if isinstance(k, str) else k for k in keys] req_id = uuid4().hex # Put an unique id in front of the request, separated by an empty # frame. This way the unique id look like an additional routing # identity to ZMQ and it will be treated as such by this client, too. request = ([req_id, "", lang, country or ""] + list( chain.from_iterable( (k, str(p) if p is not None else "") for k, p in keys))) self._socket.send_multipart( [part.encode(self._encoding) for part in request]) # Poll for the response, but only for so long... timeout = self._timeout translations = None while timeout > 0 and translations is None: start = time() sockets = dict(self._poller.poll(timeout)) timeout -= (time() - start) if self._socket in sockets: response = self._socket.recv_multipart() translations = self._handle_response(req_id, response) if translations is None: # timeout, return keys translations = [key for key, _ in keys] return translations if len(keys) > 1 else translations[0] def close(self): if self._context is not None: self._socket.close() self._socket = None self._context.term() self._context = None def translate(self, language, country, *keys): """ Translate strings with or without their plural flavor. :param language: str :param country: str :type keys: [str or (str, int)] :rtype: str or [str] """ if self._skip_translations: translations = [k if isinstance(k, str) else k[0] for k in keys] return translations if len(translations) > 1 else translations[0] if not self._translation_cache: return self.get(language, country, *keys) translations = self._translation_cache.get_multiple( language, country, keys) if all(translations): return translations if len(keys) > 1 else translations[0] # Give all the to be translated things an ordering index ordered_translations = list( zip(range(len(translations)), keys, translations)) # Split of the keys that do not have a translation indexed_missing_keys = [ (idx, key) for idx, key, translation in ordered_translations if translation is None ] idx = [idx for idx, key in indexed_missing_keys] new_keys = [key for idx, key in indexed_missing_keys] missing_translations = self.get(language, country, *new_keys) if not isinstance(missing_translations, (list, tuple)): missing_translations = [missing_translations] self._translation_cache.set_multiple(language, country, new_keys, missing_translations) missing_translations = list(zip(idx, missing_translations)) present_translations = [ (idx, translation) for idx, key, translation in ordered_translations ] final_translations = dict(present_translations) final_translations.update(dict(missing_translations)) final_translations = sorted(final_translations.items(), key=lambda x: x[0]) translations = [e[1] for e in final_translations] return translations if len(keys) > 1 else translations[0] def translate_closure(self, language, country): """ Returns a function that encloses the language and the country for this particular translation client """ def _t(*keys): return self.translate(language, country, *keys) return _t