refactor into package
This commit is contained in:
parent
ba0e9810a8
commit
05e8304a83
|
@ -0,0 +1 @@
|
||||||
|
__version__ = "0.0.1"
|
|
@ -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)
|
|
@ -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,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()
|
|
@ -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…
Reference in New Issue