├── examples
    ├── perf
    │   ├── requirements.txt
    │   ├── pub.py
    │   └── sub.py
    ├── sub.py
    └── pub.py
├── setup.py
├── README.md
└── pubsub
    ├── __init__.py
    ├── common.py
    ├── server.py
    └── client.py


/examples/perf/requirements.txt:
--------------------------------------------------------------------------------
1 | tqdm
2 | 


--------------------------------------------------------------------------------
/examples/perf/pub.py:
--------------------------------------------------------------------------------
 1 | from pubsub import MessageQueue
 2 | 
 3 | 
 4 | def main():
 5 |     m = MessageQueue()
 6 |     m.connect()
 7 |     while True:
 8 |         m.send_message('random', 'My message.')
 9 | 
10 | 
11 | if __name__ == '__main__':
12 |     main()
13 | 


--------------------------------------------------------------------------------
/examples/perf/sub.py:
--------------------------------------------------------------------------------
 1 | import sys
 2 | 
 3 | from tqdm import tqdm
 4 | 
 5 | from pubsub import MessageQueue
 6 | 
 7 | m = MessageQueue()
 8 | m.connect()
 9 | m.subscribe('random')
10 | 
11 | for i in tqdm(range(int(1e6)), desc='Benchmarking', file=sys.stdout, unit_scale=True, unit=' msg'):
12 |     m.get_message(timeout=None)
13 | 


--------------------------------------------------------------------------------
/setup.py:
--------------------------------------------------------------------------------
 1 | from distutils.core import setup
 2 | 
 3 | from setuptools import find_packages
 4 | 
 5 | setup(
 6 |     name='python-pubsub',
 7 |     version='1.4.0',
 8 |     author='Philippe Remy',
 9 |     packages=find_packages(),
10 |     install_requires=[],
11 |     entry_points={
12 |         'console_scripts': [
13 |             'start_pubsub_broker=pubsub.server:start',
14 |         ],
15 |     }
16 | )
17 | 


--------------------------------------------------------------------------------
/examples/sub.py:
--------------------------------------------------------------------------------
 1 | from pubsub import MessageQueue
 2 | 
 3 | # Start the broker with the command: start_pubsub_broker
 4 | 
 5 | m = MessageQueue()
 6 | m.connect()
 7 | count = 0
 8 | m.subscribe('prime')
 9 | m.subscribe('random')
10 | while True:
11 |     channel, message = m.get_message(timeout=None)
12 | 
13 |     if channel in ['prime', 'random']:
14 |         count += 1
15 |         print(f'RECEIVE {count}: channel: {channel}, message: {message["coremq_string"]}')
16 | 


--------------------------------------------------------------------------------
/README.md:
--------------------------------------------------------------------------------
 1 | # Python PubSub
 2 | A simple python implementation of a message router with many subscribers and many publishers.
 3 | 
 4 | It can be considered as as fork of the project: [CoreMQ](https://github.com/deejross/coremq).
 5 | 
 6 | <center>
 7 | <div><a href='//sketchviz.com/@philipperemy/8f597379f217ca0b6a20700ebc7f0b22'><img src='https://sketchviz.com/@philipperemy/8f597379f217ca0b6a20700ebc7f0b22/765f0d983410a72083cd967e7e0535dee8f9bcfc.sketchy.png' style='max-width: 100%;'></a></div>
 8 | </center>
 9 | 
10 | This implementation handles:
11 | - 26K messages per second.
12 | - Latency of around 1ms.
13 | - `TCP_NODELAY` flag activated for latency sensitive applications.
14 | - Many publishers.
15 | - Many subscribers.
16 | - Multiple channels.
17 | 
18 | ```
19 | pip install python-pubsub
20 | ```
21 | 
22 | Refer to the [examples](examples).
23 | 
24 | 


--------------------------------------------------------------------------------
/examples/pub.py:
--------------------------------------------------------------------------------
 1 | import random
 2 | 
 3 | from pubsub import MessageQueue
 4 | from time import sleep
 5 | # Start the broker with the command: start_pubsub_broker
 6 | 
 7 | 
 8 | def prime_numbers():
 9 |     for num in range(1, 1001):
10 |         for i in range(2, num):
11 |             if num % i == 0:
12 |                 break
13 |         else:
14 |             yield num
15 | 
16 | 
17 | def main():
18 |     m = MessageQueue()
19 |     m.connect()
20 |     count = 0
21 |     print(m.welcome_message)
22 |     pub_id = random.randint(0, 1000)
23 |     print(f'Publisher ID: {pub_id}.')
24 | 
25 |     def publish(channel, msg):
26 |         nonlocal count
27 |         count += 1
28 |         msg += f' Publisher Id is [{pub_id}].'
29 |         print(f'PUBLISH {count}: {msg}.')
30 |         sleep(0.0001)
31 |         m.send_message(channel, msg)
32 | 
33 |     for prime in prime_numbers():
34 |         publish('prime', f'Next prime number is [{prime}].')
35 |         publish('random', f'Next random number is [{random.randint(0, 10000)}].')
36 | 
37 | 
38 | if __name__ == '__main__':
39 |     main()
40 | 


--------------------------------------------------------------------------------
/pubsub/__init__.py:
--------------------------------------------------------------------------------
 1 | """
 2 | CoreMQ
 3 | ------
 4 | A pure-Python messaging queue.
 5 | 
 6 | License
 7 | -------
 8 | The MIT License (MIT)
 9 | Copyright (c) 2015 Ross Peoples <ross.peoples@gmail.com>
10 | Permission is hereby granted, free of charge, to any person obtaining a copy
11 | of this software and associated documentation files (the "Software"), to deal
12 | in the Software without restriction, including without limitation the rights
13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
14 | copies of the Software, and to permit persons to whom the Software is
15 | furnished to do so, subject to the following conditions:
16 | The above copyright notice and this permission notice shall be included in all
17 | copies or substantial portions of the Software.
18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
24 | SOFTWARE.
25 | """
26 | from .client import MessageQueue
27 | 


--------------------------------------------------------------------------------
/pubsub/common.py:
--------------------------------------------------------------------------------
  1 | """
  2 | CoreMQ
  3 | ------
  4 | A pure-Python messaging queue.
  5 | 
  6 | License
  7 | -------
  8 | The MIT License (MIT)
  9 | Copyright (c) 2015 Ross Peoples <ross.peoples@gmail.com>
 10 | Permission is hereby granted, free of charge, to any person obtaining a copy
 11 | of this software and associated documentation files (the "Software"), to deal
 12 | in the Software without restriction, including without limitation the rights
 13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 14 | copies of the Software, and to permit persons to whom the Software is
 15 | furnished to do so, subject to the following conditions:
 16 | The above copyright notice and this permission notice shall be included in all
 17 | copies or substantial portions of the Software.
 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 24 | SOFTWARE.
 25 | """
 26 | 
 27 | import json
 28 | import logging
 29 | import os
 30 | 
 31 | str_type = str
 32 | from configparser import ConfigParser, NoOptionError, NoSectionError
 33 | 
 34 | loggers = dict()
 35 | 
 36 | 
 37 | class ConnectionClosed(Exception):
 38 |     pass
 39 | 
 40 | 
 41 | class ProtocolError(Exception):
 42 |     pass
 43 | 
 44 | 
 45 | class CoreConfigParser(ConfigParser, object):
 46 |     def get(self, section, option, default=None):
 47 |         try:
 48 |             return super(CoreConfigParser, self).get(section, option)
 49 |         except (NoOptionError, NoSectionError):
 50 |             return default
 51 | 
 52 | 
 53 | def construct_message(queue, message):
 54 |     if not isinstance(queue, str_type):
 55 |         raise ValueError('Queue name must be a string, not %s' % queue)
 56 | 
 57 |     if len(queue) < 1:
 58 |         raise ValueError('Queue name must be at least one character in length')
 59 | 
 60 |     if ' ' in queue:
 61 |         raise ValueError('Queue name must not contain spaces')
 62 | 
 63 |     if isinstance(message, str_type):
 64 |         message = dict(coremq_string=message)
 65 | 
 66 |     if isinstance(message, dict):
 67 |         message = json.dumps(message)
 68 |     else:
 69 |         raise ValueError('Messages should be either a dictionary or a string')
 70 | 
 71 |     if len(message) > 99999999:  # 100 MB max int that can fit in message header (8 characters, plus two controls)
 72 |         raise ValueError('Message cannot be 100MB or larger')
 73 | 
 74 |     if not isinstance(message, bytes):
 75 |         message = message.encode('utf-8')
 76 | 
 77 |     return ('+%s %s ' % (len(message) + len(queue) + 1, queue)).encode('utf-8') + message
 78 | 
 79 | 
 80 | def send_message(socket, queue, message):
 81 |     socket.send(construct_message(queue, message))
 82 | 
 83 | 
 84 | def get_message(socket, timeout=1):
 85 |     socket.settimeout(timeout)
 86 |     data = socket.recv(10).decode('utf-8')
 87 |     expected_length, data = validate_header(data)
 88 | 
 89 |     while len(data) < expected_length:
 90 |         data += socket.recv(expected_length - len(data)).decode('utf-8')
 91 | 
 92 |     if ' ' not in data:
 93 |         return None, data
 94 | 
 95 |     queue, message = data.split(' ', 1)
 96 |     return queue, json.loads(message)
 97 | 
 98 | 
 99 | def validate_header(data):
100 |     """
101 |     Validates that data is in the form of "+5 Hello", with + beginning messages, followed by the length of the
102 |     message as an integer, followed by a space, then the message.
103 |     :param data: The raw data from the socket
104 |     :return: (int, str) - the expected length of the message, the message
105 |     """
106 |     if not data:
107 |         raise ConnectionClosed()
108 | 
109 |     if data[0] != '+':
110 |         raise ProtocolError('Missing beginning +')
111 | 
112 |     if ' ' not in data:
113 |         raise ProtocolError('Missing space after length')
114 | 
115 |     length, data = data.split(' ', 1)
116 | 
117 |     try:
118 |         length = int(length[1:])
119 |     except ValueError:
120 |         raise ProtocolError('Length integer must be between + and space')
121 | 
122 |     return length, data
123 | 
124 | 
125 | def load_configuration(path=None):
126 |     """
127 |     Loads configuration for CoreMQ and CoreWS servers
128 |     :param path: Optional path for the config file. Defaults to current directory
129 |     :return: CoreConfigParser
130 |     """
131 |     if not path:
132 |         path = os.path.join(os.getcwd(), 'coremq.conf')
133 | 
134 |     c = CoreConfigParser()
135 | 
136 |     try:
137 |         c.read(path)
138 |     except:
139 |         logging.error('Config file could not be loaded')
140 | 
141 |     return c
142 | 
143 | 
144 | def get_logger(config, section, logger_name=None):
145 |     """
146 |     Configures the Python logging module based on the configuration settings from load_configuration
147 |     :param config:  The CoreConfigParser instance from load_configuration
148 |     :param section: The section where the logging config can be found
149 |     :param logger_name: The name of the logger. Uses section by default
150 |     :return: Logger
151 |     """
152 |     logging.basicConfig()
153 | 
154 |     if section in loggers:
155 |         return loggers[section]
156 | 
157 |     logger = logging.getLogger(section or logger_name)
158 |     logger.propagate = False
159 | 
160 |     log_file = config.get('CoreMQ', 'log_file', 'stdout')
161 |     if log_file == 'stdout':
162 |         handler = logging.StreamHandler()
163 |     else:
164 |         handler = logging.FileHandler(log_file)
165 | 
166 |     log_level = config.get('CoreMQ', 'log_level', 'DEBUG')
167 |     logger.setLevel(logging.getLevelName(log_level))
168 | 
169 |     formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
170 |     handler.setFormatter(formatter)
171 | 
172 |     logger.addHandler(handler)
173 |     loggers[section] = logger
174 |     return logger
175 | 
176 | 
177 | def comma_string_to_list(s):
178 |     """
179 |     Takes a line of values, seprated by commas and returns the values in a list, removing any extra whitespacing.
180 |     :param s: The string with commas
181 |     :return: list
182 |     """
183 |     if isinstance(s, (list, tuple)):
184 |         return s
185 | 
186 |     result = []
187 |     items = s.split(',')
188 |     for i in items:
189 |         result.append(i.strip())
190 | 
191 |     return result
192 | 


--------------------------------------------------------------------------------
/pubsub/server.py:
--------------------------------------------------------------------------------
  1 | """
  2 | CoreMQ
  3 | ------
  4 | A pure-Python messaging queue.
  5 | 
  6 | License
  7 | -------
  8 | The MIT License (MIT)
  9 | Copyright (c) 2015 Ross Peoples <ross.peoples@gmail.com>
 10 | Permission is hereby granted, free of charge, to any person obtaining a copy
 11 | of this software and associated documentation files (the "Software"), to deal
 12 | in the Software without restriction, including without limitation the rights
 13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 14 | copies of the Software, and to permit persons to whom the Software is
 15 | furnished to do so, subject to the following conditions:
 16 | The above copyright notice and this permission notice shall be included in all
 17 | copies or substantial portions of the Software.
 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 24 | SOFTWARE.
 25 | """
 26 | import collections
 27 | import signal
 28 | import socket
 29 | import sys
 30 | import time
 31 | import traceback
 32 | import uuid
 33 | from multiprocessing import Process
 34 | from socketserver import BaseRequestHandler, ThreadingMixIn, TCPServer
 35 | 
 36 | from pubsub.common import ConnectionClosed, get_message, send_message
 37 | 
 38 | ADDRESS = '127.0.0.1'
 39 | PORT = 6747  # spells MSGQ (message queue)
 40 | 
 41 | 
 42 | class ThreadedTCPServer(ThreadingMixIn, TCPServer):
 43 |     EXITING = False
 44 |     PROCESS = None
 45 |     HISTORY = dict()
 46 | 
 47 | 
 48 | ThreadedTCPServer.allow_reuse_address = True
 49 | 
 50 | 
 51 | class TCPRequestHandler(BaseRequestHandler):
 52 |     connections = dict()
 53 | 
 54 |     def handle(self):
 55 |         conn_id = str(uuid.uuid4())
 56 |         TCPRequestHandler.connections[conn_id] = dict(handler=self, subscriptions=[conn_id], options=dict())
 57 |         self.respond(conn_id, 'Welcome!')
 58 | 
 59 |         print('Clients connected: %s' % len(TCPRequestHandler.connections))
 60 |         while True:
 61 |             try:
 62 |                 queue, message = get_message(self.request)
 63 |                 message['coremq_sender'] = conn_id
 64 |                 message['coremq_sent'] = time.time()
 65 | 
 66 |                 if 'coremq_subscribe' in message:
 67 |                     self.subscribe(conn_id, message['coremq_subscribe'])
 68 |                     self.respond(conn_id, 'OK: Subscribe successful')
 69 |                 elif 'coremq_unsubscribe' in message:
 70 |                     self.unsubscribe(conn_id, message['coremq_unsubscribe'])
 71 |                     self.respond(conn_id, 'OK: Unsubscribe successful')
 72 |                 elif 'coremq_options' in message:
 73 |                     self.set_options(conn_id, message['coremq_options'])
 74 |                     self.respond(conn_id, 'OK: Options set')
 75 |                 elif 'coremq_gethistory' in message:
 76 |                     self.get_history(conn_id, message['coremq_gethistory'])
 77 |                 else:
 78 |                     self.respond(conn_id, 'OK: Message sent')
 79 |                     self.broadcast(queue, message)
 80 |                     self.store_message(queue, message)
 81 | 
 82 |             except socket.timeout:
 83 |                 pass
 84 |             except (ConnectionClosed, socket.error):
 85 |                 break
 86 |             except Exception as ex:
 87 |                 self.respond(conn_id, str(ex))
 88 |                 print(conn_id, ex)
 89 |                 traceback.print_exc()
 90 | 
 91 |         try:
 92 |             self.respond(conn_id, 'BYE')
 93 |         except socket.error:
 94 |             pass
 95 | 
 96 |         if conn_id in TCPRequestHandler.connections:
 97 |             del TCPRequestHandler.connections[conn_id]
 98 | 
 99 |         print('Clients connected: %s' % len(TCPRequestHandler.connections))
100 | 
101 |     def respond(self, conn_id, text):
102 |         send_message(self.request, conn_id, dict(response=text))
103 | 
104 |     def subscribe(self, conn_id, queues):
105 |         if not queues:
106 |             return
107 | 
108 |         if conn_id not in TCPRequestHandler.connections:
109 |             return
110 | 
111 |         if not isinstance(queues, (list, tuple)):
112 |             queues = [queues]
113 | 
114 |         subs = TCPRequestHandler.connections[conn_id]['subscriptions']
115 |         for q in queues:
116 |             if q not in subs:
117 |                 subs.append(q)
118 | 
119 |     def unsubscribe(self, conn_id, queues):
120 |         if not queues:
121 |             return
122 | 
123 |         if conn_id not in TCPRequestHandler.connections:
124 |             return
125 | 
126 |         if not isinstance(queues, (list, tuple)):
127 |             queues = [queues]
128 | 
129 |         subs = TCPRequestHandler.connections[conn_id]['subscriptions']
130 |         for q in queues:
131 |             if q in subs:
132 |                 subs.remove(q)
133 | 
134 |     def set_options(self, conn_id, options):
135 |         opts = TCPRequestHandler.connections[conn_id]['options']
136 |         opts.update(options)
137 | 
138 |         for key, val in options.items():
139 |             if val is None and key in opts:
140 |                 del opts[key]
141 | 
142 |     def broadcast(self, queue, message):
143 |         for conn_id, d in TCPRequestHandler.connections.items():
144 |             if conn_id == message['coremq_sender'] and queue != conn_id and d['options'].get('echo', False) is False:
145 |                 continue
146 | 
147 |             if queue in d['subscriptions']:
148 |                 send_message(d['handler'].request, queue, message)
149 | 
150 |     def store_message(self, queue, message):
151 |         if not queue in ThreadedTCPServer.HISTORY:
152 |             ThreadedTCPServer.HISTORY[queue] = collections.deque(maxlen=10)
153 | 
154 |         ThreadedTCPServer.HISTORY[queue].append(message)
155 | 
156 |     def get_history(self, conn_id, queues):
157 |         result = dict()
158 |         for q in queues:
159 |             if q in ThreadedTCPServer.HISTORY:
160 |                 result[q] = list(ThreadedTCPServer.HISTORY[q])
161 | 
162 |         send_message(self.request, conn_id, dict(response=result))
163 | 
164 | 
165 | def signal_handler(signal, frame):
166 |     print('Shutting down message queue')
167 |     ThreadedTCPServer.EXITING = True
168 | 
169 | 
170 | def message_queue_process():
171 |     print('Starting message queue')
172 |     server = ThreadedTCPServer((ADDRESS, PORT), TCPRequestHandler)
173 |     server.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
174 |     server.timeout = 1
175 |     while not ThreadedTCPServer.EXITING:
176 |         server.handle_request()
177 | 
178 |     print('Message queue stopped')
179 |     sys.exit(0)
180 | 
181 | 
182 | def start():
183 |     signal.signal(signal.SIGINT, signal_handler)
184 |     signal.signal(signal.SIGTERM, signal_handler)
185 |     ThreadedTCPServer.PROCESS = Process(target=message_queue_process)
186 |     ThreadedTCPServer.PROCESS.start()
187 | 
188 | 
189 | if __name__ == '__main__':
190 |     start()
191 | 


--------------------------------------------------------------------------------
/pubsub/client.py:
--------------------------------------------------------------------------------
  1 | """
  2 | CoreMQ
  3 | ------
  4 | A pure-Python messaging queue.
  5 | 
  6 | License
  7 | -------
  8 | The MIT License (MIT)
  9 | Copyright (c) 2015 Ross Peoples <ross.peoples@gmail.com>
 10 | Permission is hereby granted, free of charge, to any person obtaining a copy
 11 | of this software and associated documentation files (the "Software"), to deal
 12 | in the Software without restriction, including without limitation the rights
 13 | to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
 14 | copies of the Software, and to permit persons to whom the Software is
 15 | furnished to do so, subject to the following conditions:
 16 | The above copyright notice and this permission notice shall be included in all
 17 | copies or substantial portions of the Software.
 18 | THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
 19 | IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
 20 | FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
 21 | AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
 22 | LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 23 | OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 24 | SOFTWARE.
 25 | """
 26 | 
 27 | import socket
 28 | from json import JSONDecodeError
 29 | 
 30 | from pubsub.common import get_message, send_message, ProtocolError
 31 | 
 32 | 
 33 | class MessageQueue(object):
 34 |     def __init__(self, server='127.0.0.1', port=6747):
 35 |         self.server = server
 36 |         self.port = port
 37 |         self.socket = None
 38 |         self.connection_id = None
 39 |         self.welcome_message = None
 40 |         self.subscriptions = []
 41 |         self.options = dict()
 42 |         self.last_message_time = 0
 43 | 
 44 |     def connect(self):
 45 |         if self.socket:
 46 |             return
 47 | 
 48 |         self.socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 49 |         self.socket.settimeout(30)
 50 |         self.socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
 51 | 
 52 |         self.socket.connect((self.server, self.port))
 53 |         self.connection_id, self.welcome_message = get_message(self.socket)
 54 | 
 55 |         if self.subscriptions:
 56 |             self.subscribe(*self.subscriptions)
 57 | 
 58 |         if self.options:
 59 |             self.set_options(**self.options)
 60 | 
 61 |     def close(self):
 62 |         if self.socket:
 63 |             self.socket.close()
 64 |             self.socket = None
 65 | 
 66 |     def send_message(self, queue, message):
 67 |         if not self.socket:
 68 |             self.connect()
 69 | 
 70 |         try:
 71 |             send_message(self.socket, queue, message)
 72 |         except socket.error:
 73 |             # attempt to reconnect if there was a connection error
 74 |             self.close()
 75 |             self.connect()
 76 |             send_message(self.socket, queue, message)
 77 | 
 78 |         # try:
 79 |         #     return self.get_message()
 80 |         # except socket.error:
 81 |         #     return None, None
 82 | 
 83 |     def get_message(self, timeout=1):
 84 |         if not self.socket:
 85 |             self.connect()
 86 | 
 87 |         try:
 88 |             queue, message = get_message(self.socket, timeout=timeout)
 89 |         except socket.timeout:
 90 |             return None, None
 91 |         except socket.error:
 92 |             # attempt to reconnect if there was a connection error
 93 |             self.close()
 94 |             self.connect()
 95 |             try:
 96 |                 queue, message = get_message(self.socket, timeout=timeout)
 97 |             except socket.timeout:
 98 |                 return None, None
 99 | 
100 |         if 'response' in message and message['response'] == 'BYE':
101 |             self.close()
102 | 
103 |         return queue, message
104 | 
105 |     def listen(self, seconds=30):
106 |         for i in range(seconds):
107 |             m = self.get_message()
108 |             if m[0]:
109 |                 print(m[0], m[1])
110 | 
111 |     def stress(self, queue, count=10000, wait=0, silent=False):
112 |         import time
113 | 
114 |         start_time = time.time()
115 |         for i in range(count):
116 |             self.send_message(queue, dict(iteration=i))
117 |             if wait:
118 |                 time.sleep(wait)
119 |         end_time = time.time()
120 | 
121 |         if not silent:
122 |             print('Stress results:')
123 |             print('Iterations: %s' % count)
124 |             print('Forced wait time: %s' % wait)
125 |             print('Time taken: %s seconds' % (end_time - start_time))
126 |             print('Messages per second: %s' % (float(count) / (end_time - start_time)))
127 | 
128 |     def get_history(self, *queues):
129 |         if not queues:
130 |             queues = self.subscriptions
131 | 
132 |         if not queues:
133 |             raise ValueError('Must pass at least one queue name')
134 | 
135 |         if not isinstance(queues, (list, tuple)):
136 |             queues = [queues]
137 | 
138 |         return self.send_message(self.connection_id, dict(coremq_gethistory=queues))
139 | 
140 |     def subscribe(self, *queues):
141 |         if not queues:
142 |             raise ValueError('Must pass at least one queue name')
143 | 
144 |         if not isinstance(queues, (list, tuple)):
145 |             queues = [queues]
146 | 
147 |         for q in queues:
148 |             if q not in self.subscriptions:
149 |                 self.subscriptions.append(q)
150 | 
151 |         return self.send_message(self.connection_id, dict(coremq_subscribe=queues))
152 | 
153 |     def unsubscribe(self, *queues):
154 |         if not queues:
155 |             raise ValueError('Must pass at least one queue name')
156 | 
157 |         if not isinstance(queues, (list, tuple)):
158 |             queues = [queues]
159 | 
160 |         for q in queues:
161 |             if q in self.subscriptions:
162 |                 self.subscriptions.remove(q)
163 | 
164 |         return self.send_message(self.connection_id, dict(coremq_unsubscribe=queues))
165 | 
166 |     def set_options(self, **options):
167 |         self.options.update(options)
168 | 
169 |         for key, val in options.items():
170 |             if val is None and key in self.options:
171 |                 del self.options[key]
172 | 
173 |         return self.send_message(self.connection_id, dict(coremq_options=options))
174 | 
175 | 
176 | def stress_worker(args):
177 |     import random
178 |     import time
179 |     server, port, queue, count, wait, listen_only = args
180 | 
181 |     time.sleep(random.uniform(0, 1))
182 |     m = MessageQueue(server, port)
183 |     try:
184 |         m.connect()
185 |     except:
186 |         time.sleep(1)
187 |         m.connect()
188 | 
189 |     if listen_only:
190 |         try:
191 |             m.subscribe(queue)
192 |             messages = 0
193 |             null_msg = 0
194 |             errors = 0
195 |             for i in range(count):
196 |                 try:
197 |                     msg = m.get_message(timeout=10)
198 |                     if msg[0] is not None:
199 |                         messages += 1
200 |                     else:
201 |                         null_msg += 1
202 |                         if null_msg >= 2:
203 |                             break
204 |                 except:
205 |                     errors += 1
206 |         except:
207 |             return count, 1
208 | 
209 |         m.close()
210 | 
211 |         return count - messages, errors
212 |     else:
213 |         time.sleep(random.uniform(0, 1))
214 |         errors = 0
215 |         start_time = time.time()
216 |         for i in range(count):
217 |             try:
218 |                 m.send_message(queue, dict(iteration=i + 1))
219 |                 if wait:
220 |                     time.sleep(wait)
221 |             except:
222 |                 errors += 1
223 | 
224 |         end_time = time.time()
225 | 
226 |         m.close()
227 |         return end_time - start_time, float(count) / (end_time - start_time), errors
228 | 
229 | 
230 | def parallel_stress(server, port, queue, count=1000, wait=0.01, workers=10, listeners=90):
231 |     import multiprocessing
232 | 
233 |     print('Beginning stress test with these settings:')
234 |     print('Server: %s' % server)
235 |     print('Queue: %s' % queue)
236 |     print('Count per worker: %s' % count)
237 |     print('Delay between messages: %s' % wait)
238 |     print('Number of workers: %s' % workers)
239 |     print('Number of listeners: %s' % listeners)
240 |     print('')
241 | 
242 |     pool1 = multiprocessing.Pool(listeners)
243 |     pool2 = multiprocessing.Pool(workers)
244 |     tasks1 = [(server, port, queue, count, wait, True) for i in range(listeners)]
245 |     tasks2 = [(server, port, queue, count, wait, False) for i in range(workers)]
246 | 
247 |     job1 = pool1.map_async(stress_worker, tasks1)
248 |     job2 = pool2.map_async(stress_worker, tasks2)
249 |     job1.wait()
250 |     job2.wait()
251 |     listen_results = job1.get()
252 |     worker_results = job2.get()
253 |     slowest = 0
254 |     fastest = 9999
255 |     average_mps = 0
256 |     errors = 0
257 |     missed = 0
258 | 
259 |     for t, mps, errs in worker_results:
260 |         if t > slowest:
261 |             slowest = t
262 | 
263 |         if t < fastest:
264 |             fastest = t
265 | 
266 |         average_mps += mps
267 |         errors += errs
268 | 
269 |     average_mps /= float(workers)
270 | 
271 |     for m, errs in listen_results:
272 |         missed += m
273 | 
274 |     print('Test results:')
275 |     print('Total count: %s' % (count * workers))
276 |     print('Fastest worker: %s seconds' % fastest)
277 |     print('Slowest worker: %s seconds' % slowest)
278 |     print('Number of errors: %s' % errors)
279 |     print('Number of missing messages: %s' % missed)
280 |     print('Average messages per second: %s' % average_mps)
281 | 


--------------------------------------------------------------------------------