class Discover(object):
def __init__(self, port=8080, *args, **kwargs):
self._ctx = zmq.Context()
self._terminated = False # API shut us down
self._verbose = False # Log all traffic (logging module?)
self.beacon_port = ZRE_DISCOVERY_PORT # Beacon port number
self.interval = 0 # Beacon interval 0=default
self.beacon = None # Beacon actor
self.beacon_socket = None # Beacon socket for polling
self.poller = zmq.Poller() # Socket poller
self.identity = uuid.uuid4() # Our UUID as object
self.bound = False
self.name = str(self.identity)[:6] # Our public name (default=first 6 uuid chars)
self.endpoint = "" # Our public endpoint
self.port = port # Our inbox port, if any
self.status = 0 # Our own change counter
self.peers = {} # Hash of known peers, fast lookup
self.headers = {} # Our header values
# TODO: gossip stuff
# self.start()
self.run_thread = Thread(target=self.run)
self.run_thread.start()
# def __del__(self):
# destroy beacon
def start(self):
# TODO: If application didn't bind explicitly, we grab an ephemeral port
# on all available network interfaces. This is orthogonal to
# beaconing, since we can connect to other peers and they will
# gossip our endpoint to others.
if self.beacon_port:
# Start beacon discovery
self.beacon = ZActor(self._ctx, ZBeacon)
if self._verbose:
self.beacon.send_unicode("VERBOSE")
# Our hostname is provided by zbeacon
self.beacon.send_unicode("CONFIGURE", zmq.SNDMORE)
self.beacon.send(struct.pack("I", self.beacon_port))
hostname = self.beacon.recv_unicode()
self.endpoint = "tcp://%s:%d" % (hostname, self.port)
# Set broadcast/listen beacon
transmit = struct.pack('cccb16sH', b'Z', b'R', b'E',
BEACON_VERSION, self.identity.bytes,
socket.htons(self.port))
self.beacon.send_unicode("PUBLISH", zmq.SNDMORE)
self.beacon.send(transmit)
# construct the header filter (to discard none zre messages)
filter = struct.pack("ccc", b'Z', b'R', b'E')
self.beacon.send_unicode("SUBSCRIBE", zmq.SNDMORE)
self.beacon.send(filter)
self.beacon_socket = self.beacon.resolve()
self.poller.register(self.beacon_socket, zmq.POLLIN)
def stop(self):
logger.debug("Pyre node: stopping beacon")
if self.beacon:
stop_transmit = struct.pack('cccb16sH', b'Z', b'R', b'E',
BEACON_VERSION, self.identity.bytes,
socket.htons(0))
self.beacon.send_unicode("PUBLISH", zmq.SNDMORE)
self.beacon.send(stop_transmit)
# Give time for beacon to go out
time.sleep(0.001)
self._terminated = True
# Give time for thread to exit
time.sleep(5)
self.poller.unregister(self.beacon_socket)
#self.beacon.send_unicode("$TERM")
self.beacon.destroy()
self.beacon = None
self.beacon_socket = None
self.beacon_port = 0
# Find or create peer via its UUID string
def require_peer(self, identity, endpoint):
"""
:param identity:
:param endpoint:
:return: endpoint of peer
"""
p = self.peers.get(identity)
if not p:
self.peers[identity] = endpoint
return p
def get_peers(self):
return list(self.peers.values())
# Remove a peer from our data structures
def remove_peer(self, peer):
# To destroy peer, we remove from peers hash table (dict)
self.peers.pop(peer)
def recv_beacon(self):
# Get IP address and beacon of peer
msgs = self.beacon_socket.recv_multipart()
ipaddress = msgs.pop(0)
frame = msgs.pop(0)
beacon = struct.unpack('cccb16sH', frame)
# Ignore anything that isn't a valid beacon
if beacon[3] != BEACON_VERSION:
logger.warning("Invalid ZRE Beacon version: {0}".format(beacon[3]))
return
peer_id = uuid.UUID(bytes=beacon[4])
port = socket.ntohs(beacon[5])
# if we receive a beacon with port 0 this means the peer exited
if port:
peer = self.require_peer(peer_id, (str(ipaddress.decode('UTF-8')), port))
else:
# Zero port means peer is going away; remove it if
# we had any knowledge of it already
peer = self.peers.get(peer_id)
# remove the peer (delete)
if peer:
logger.debug("Received 0 port beacon, removing peer {0}".format(peer))
self.remove_peer(peer_id)
else:
logger.warning(self.peers)
logger.warning("We don't know peer id {0}".format(peer_id))
# TODO: Handle gossip dat
# --------------------------------------------------------------------------
# This is the actor that runs a single node; it uses one thread, creates
# a zyre_node object at start and destroys that when finishing.
def run(self):
while not self._terminated:
items = dict(self.poller.poll(1000))
if self.beacon_socket in items and items[self.beacon_socket] == zmq.POLLIN:
self.recv_beacon()
# Poll on API pipe and on UDP socket
items = dict(self.poller.poll(timeout * 1000))
if self.pipe in items and items[self.pipe] == zmq.POLLIN:
self.handle_pipe()
if self.udpsock.fileno() in items and items[self.udpsock.fileno()] == zmq.POLLIN:
self.handle_udp()
if self.transmit and time.time() >= self.ping_at:
self.send_beacon()
self.ping_at = time.time() + self.interval
if self.terminated:
break
if __name__ == '__main__':
import zmq
import struct
import time
speaker = ZActor(zmq.Context(), ZBeacon)
speaker.send_unicode("VERBOSE")
speaker.send_unicode("CONFIGURE", zmq.SNDMORE)
speaker.send(struct.pack("I", 9999))
speaker.send_unicode("PUBLISH", zmq.SNDMORE)
import uuid
transmit = struct.pack('cccb16sH', b'Z', b'R', b'E',
1, uuid.uuid4().bytes,
socket.htons(1300))
speaker.send(transmit)
speaker.destroy()
class ZBeaconTest(unittest.TestCase):
def setUp(self, *args, **kwargs):
ctx = zmq.Context()
ctx = zmq.Context()
# two beacon frames
self.transmit1 = struct.pack("cccb16sH", b"Z", b"R", b"E", 1, uuid.uuid4().bytes, socket.htons(9999))
self.transmit2 = struct.pack("cccb16sH", b"Z", b"R", b"E", 1, uuid.uuid4().bytes, socket.htons(9999))
self.node1 = ZActor(ctx, ZBeacon)
self.node1.send_unicode("VERBOSE")
self.node1.send_unicode("CONFIGURE", zmq.SNDMORE)
self.node1.send(struct.pack("I", 9999))
print("Hostname 1:", self.node1.recv_unicode())
self.node2 = ZActor(ctx, ZBeacon)
self.node2.send_unicode("VERBOSE")
self.node2.send_unicode("CONFIGURE", zmq.SNDMORE)
self.node2.send(struct.pack("I", 9999))
print("Hostname 2:", self.node2.recv_unicode())
# end setUp
def tearDown(self):
self.node1.destroy()
self.node2.destroy()
# end tearDown
def test_node1(self):
self.node1.send_unicode("PUBLISH", zmq.SNDMORE)
self.node1.send(self.transmit1)
def test_node2(self):
self.node2.send_unicode("PUBLISH", zmq.SNDMORE)
self.node2.send(self.transmit2)
def test_recv_beacon1(self):
self.node1.send_unicode("PUBLISH", zmq.SNDMORE)
self.node1.send(self.transmit1)
self.node2.send_unicode("PUBLISH", zmq.SNDMORE)
self.node2.send(self.transmit2)
req = self.node1.recv_multipart()
self.assertEqual(self.transmit2, req[1])
def test_recv_beacon1(self):
self.node1.send_unicode("PUBLISH", zmq.SNDMORE)
self.node1.send(self.transmit1)
self.node2.send_unicode("PUBLISH", zmq.SNDMORE)
self.node2.send(self.transmit2)
req = self.node2.recv_multipart()
self.assertEqual(self.transmit1, req[1])
class ZBeaconTest(unittest.TestCase):
def setUp(self, *args, **kwargs):
ctx = zmq.Context()
ctx = zmq.Context()
# two beacon frames
self.transmit1 = struct.pack('cccb16sH', b'Z', b'R', b'E',
1, uuid.uuid4().bytes,
socket.htons(9999))
self.transmit2 = struct.pack('cccb16sH', b'Z', b'R', b'E',
1, uuid.uuid4().bytes,
socket.htons(9999))
self.node1 = ZActor(ctx, ZBeacon)
self.node1.send_unicode("VERBOSE")
self.node1.send_unicode("CONFIGURE", zmq.SNDMORE)
self.node1.send(struct.pack("I", 9999))
print("Hostname 1:", self.node1.recv_unicode())
self.node2 = ZActor(ctx, ZBeacon)
self.node2.send_unicode("VERBOSE")
self.node2.send_unicode("CONFIGURE", zmq.SNDMORE)
self.node2.send(struct.pack("I", 9999))
print("Hostname 2:", self.node2.recv_unicode())
# end setUp
def tearDown(self):
self.node1.destroy()
self.node2.destroy()
# end tearDown
def test_node1(self):
self.node1.send_unicode("PUBLISH", zmq.SNDMORE)
self.node1.send(self.transmit1)
def test_node2(self):
self.node2.send_unicode("PUBLISH", zmq.SNDMORE)
self.node2.send(self.transmit2)
def test_recv_beacon1(self):
self.node1.send_unicode("PUBLISH", zmq.SNDMORE)
self.node1.send(self.transmit1)
self.node2.send_unicode("PUBLISH", zmq.SNDMORE)
self.node2.send(self.transmit2)
req = self.node1.recv_multipart()
self.assertEqual(self.transmit2, req[1])
def test_recv_beacon1(self):
self.node1.send_unicode("PUBLISH", zmq.SNDMORE)
self.node1.send(self.transmit1)
self.node2.send_unicode("PUBLISH", zmq.SNDMORE)
self.node2.send(self.transmit2)
req = self.node2.recv_multipart()
self.assertEqual(self.transmit1, req[1])
class ZBeaconRepeater(object):
def __init__(self):
self.dont_ip = os.environ.get('DONT_IP', '').split(' ')
self.other_repeaters = []
self.consul = Consul(sys.argv[1])
self.ctx = zmq.Context()
self.poller = zmq.Poller()
self.repeater_pub_port = REPEATER_PUB_PORT
self.pub = self.ctx.socket(zmq.PUB)
self.pub.bind('tcp://*:%d' % self.repeater_pub_port)
self.sub = self.ctx.socket(zmq.SUB)
self.sub.setsockopt(zmq.SUBSCRIBE, b'')
self.poller.register(self.sub, zmq.POLLIN)
self.beacon = ZActor(self.ctx, ZBeacon)
self.beacon.send_unicode('CONFIGURE', zmq.SNDMORE)
self.beacon.send(struct.pack('I', ZRE_DISCOVERY_PORT))
self.address = self.beacon.recv_unicode() # Hostname
filter_ = struct.pack('ccc', b'Z', b'R', b'E')
self.beacon.send_unicode('SUBSCRIBE',zmq.SNDMORE)
self.beacon.send(filter_)
self.beacon_socket = self.beacon.resolve()
self.poller.register(self.beacon_socket, zmq.POLLIN)
def _connect_to_repeaters(self):
logger.debug('Syncing repeaters')
docker_nodes = self.consul.kv.get('docker/nodes', recurse=True)
for node in docker_nodes[1]:
ip = node['Value'].split(':')[0]
if ip in self.dont_ip:
continue
self._connect_to_repeater(ip)
def _connect_to_repeater(self, ip):
if ip in self.other_repeaters:
return
logger.info('Connecting to repeater %s', ip)
endpoint = 'tcp://%s:%d' % (ip, self.repeater_pub_port)
self.sub.connect(endpoint)
self.other_repeaters.append(ip)
def run(self):
try:
time_sync_repeaters = time.time() + DELAY_SYNC_REPEATERS
while True:
if time.time() >= time_sync_repeaters:
self._connect_to_repeaters()
time_sync_repeaters = time.time() + DELAY_SYNC_REPEATERS
items = dict(self.poller.poll(0.1))
while len(items) > 0:
for fd, ev in items.items():
if (self.sub == fd) and (ev == zmq.POLLIN):
data = self.sub.recv_multipart()
logger.debug('From SUB: %s', data)
self.beacon.send_unicode('SEND BEACON', zmq.SNDMORE)
self.beacon.send(data[0])
elif (self.beacon_socket == fd) and (ev == zmq.POLLIN):
addr, data = self.beacon_socket.recv_multipart()
if addr == self.address:
continue
logger.debug('From UDP: %s', data + addr)
self.pub.send(data + addr)
items = dict(self.poller.poll(0))
finally:
self.beacon.destroy()
self.sub.close()
self.pub.close()