Browse Source

refactor into package

master
dave 4 years ago
parent
commit
05e8304a83
  1. 1
      msgbus/__init__.py
  2. 86
      msgbus/client.py
  3. 50
      msgbus/pub.py
  4. 0
      msgbus/server.py
  5. 56
      msgbus/sub.py
  6. 19
      setup.py

1
msgbus/__init__.py

@ -0,0 +1 @@
__version__ = "0.0.1"

86
msgbus/client.py

@ -0,0 +1,86 @@
import zmq
from threading import Semaphore
from time import sleep, time
class PublishSetupException(Exception):
pass
class MsgbusSubClient(object):
def __init__(self, host, port):
self.host = host
self.port = port
self.ctx = None # ZMQ context
self.sub_socket = None # listener sockets
self.subscriptions = []
self.pub_socket = None # publisher socket
self.lock = Semaphore(1)
self.connect()
def close(self):
if self.sub_socket:
self.sub_socket.close()
if self.pub_socket:
self.pub_socket.close()
if self.ctx:
self.ctx.destroy()
def connect(self):
if not self.ctx:
self.ctx = zmq.Context()
self.sub_socket = self.ctx.socket(zmq.SUB)
self.sub_socket.connect("tcp://{}:{}".format(self.host, self.port))
def sub(self, channel=None):
if channel is None:
channel = ''
assert type(channel) is str
self.sub_socket.setsockopt(zmq.SUBSCRIBE, channel.encode("utf-8"))
self.subscriptions.append(channel)
def unsub(self, channel):
assert type(channel) is str
if channel in self.subscriptions:
self.subscriptions.remove(channel)
self.sub_socket.setsockopt(zmq.UNSUBSCRIBE, channel.encode("utf-8"))
def recv(self, decode=True, block=True):
recv_args = (zmq.NOBLOCK, ) if not block else ()
message = self.sub_socket.recv(*recv_args)
channel, body = message.split(b' ', 1)
return channel.decode("utf-8"), (body.decode('utf-8') if decode else body)
def _setup_publish_socket(self, timeout=5):
start = time()
try:
self.sub("__msgbus_meta")
while not timeout or time() < start + timeout:
try:
channel, message = self.recv(block=False)
except zmq.error.Again:
sleep(0.01)
continue
if channel != "__msgbus_meta":
continue
meta, args = message.split(" ", 1)
if meta != "__my_info":
continue
server_name, subport, subproto, pubbport, pubproto = args.split(" ")
remote = "tcp://{}:{}".format(self.host, subport)
pub_socket = self.ctx.socket(zmq.PUB)
pub_socket.connect(remote)
return pub_socket
raise PublishSetupException("Could not establish publisher socket")
finally:
self.unsub("__msgbus_meta")
def pub(self, channel, message, encode_msg=True, settle=True, timeout=5):
if encode_msg:
message = message.encode("utf-8")
if not self.pub_socket:
with self.lock:
self.pub_socket = self._setup_publish_socket(timeout)
if settle:
sleep(1)
self.pub_socket.send(channel.encode("utf-8") + b' ' + message)

50
msgbus/pub.py

@ -0,0 +1,50 @@
def send_native(host, port, channel, messages):
"""
Send some messages on a specified channel using the msgbus client
"""
from contextlib import closing
from msgbus.client import MsgbusSubClient
with closing(MsgbusSubClient(host, port)) as client:
for message in messages:
client.pub(channel, message)
def send_zmq(host, port, channel, messages):
"""
Send some messages on a specified channel using a raw zmq socket.
Note: the native client connects to the server's publisher port and autodetects the server's subscriber port. The
raw zmq socket must connect directly to the subscriber port
"""
import zmq
from time import sleep
with zmq.Context() as ctx:
forward_socket = ctx.socket(zmq.PUB)
forward_socket.connect("tcp://{}:{}".format(host, port))
sleep(1)
for message in messages:
m = "{} {}".format(channel, message).encode("utf-8")
print(m)
forward_socket.send(m)
sleep(1)
forward_socket.close()
def main():
import argparse
parser = argparse.ArgumentParser(description="send a message to a msgbus server")
parser.add_argument("-i", "--host", default="127.0.0.1", help="host to connect to")
parser.add_argument("-p", "--port", default=7003, help="port to connect to")
parser.add_argument("-c", "--channel", required=True, help="message channel")
parser.add_argument("-m", "--message", required=True, nargs="+", help="message bodies")
parser.add_argument("--type", default="native", choices=["native", "raw"], help="client type")
args = parser.parse_args()
if args.type == "native":
send_native(args.host, args.port, args.channel, args.message)
elif args.type == "raw":
send_zmq(args.host, args.port, args.channel, args.message)
if __name__ == '__main__':
main()

0
server.py → msgbus/server.py

56
msgbus/sub.py

@ -0,0 +1,56 @@
def listen_native(host, port, channels):
"""
Subscribe to the given server/channels using the msgbus client
"""
from contextlib import closing
from msgbus.client import MsgbusSubClient
with closing(MsgbusSubClient(host, port)) as client:
if channels:
for channel in channels:
client.sub(channel)
else:
client.sub()
while True:
yield "{} {}".format(*client.recv())
def listen_zmq(host, port, channels):
"""
Example subscribing to the given server/channels using a raw zeromq socket
"""
import zmq
with zmq.Context() as ctx:
forward_socket = ctx.socket(zmq.SUB)
forward_socket.connect("tcp://{}:{}".format(host, port))
if channels:
for channel in channels:
forward_socket.setsockopt(zmq.SUBSCRIBE, channel.encode("utf-8"))
else:
forward_socket.setsockopt(zmq.SUBSCRIBE, b'')
while True:
transport = forward_socket.recv()
yield transport.decode("utf-8")
def main():
import argparse
parser = argparse.ArgumentParser(description="dump all messages from msgbus")
parser.add_argument("-i", "--host", default="127.0.0.1", help="host to connect to")
parser.add_argument("-p", "--port", default=7003, help="port to connect to")
parser.add_argument("-c", "--channel", nargs="+", help="dump only channels")
parser.add_argument("--type", default="native", choices=["native", "raw"], help="client type")
args = parser.parse_args()
client = None
if args.type == "native":
client = listen_native(args.host, args.port, args.channel)
elif args.type == "raw":
client = listen_zmq(args.host, args.port, args.channel)
for line in client:
print(line)
if __name__ == '__main__':
main()

19
setup.py

@ -0,0 +1,19 @@
#!/usr/bin/env python3
from setuptools import setup
from msgbus import __version__
setup(name='msgbus',
version=__version__,
description='pubsub communication tools',
url='http://gitlab.xmopx.net/dave/pymsgbus',
author='dpedu',
author_email='dave@davepedu.com',
packages=['msgbus'],
entry_points={
"console_scripts": [
"mbusd = msgbus.server:main",
"mbuspub = msgbus.pub:main",
"mbussub = msgbus.sub:main",
]
})
Loading…
Cancel
Save