occheung
/
artiq
forked from M-Labs/artiq
1
0
Fork 0
Code Pull Requests Activity

protocols/broadcast,sync_struct: add disconnect callback

This commit is contained in:
Sebastien Bourdeauducq 2017-01-30 10:47:56 +08:00
parent 43aad0914e
commit dcea48a52d
2 changed files with 36 additions and 22 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

8
artiq/protocols/broadcast.py
View File

@ -9,11 +9,12 @@ _init_string = b"ARTIQ broadcast\n"
class Receiver: class Receiver:
def __init__(self, name, notify_cb): def __init__(self, name, notify_cb, disconnect_cb=None):
self.name = name self.name = name
if not isinstance(notify_cb, list): if not isinstance(notify_cb, list):
notify_cb = [notify_cb] notify_cb = [notify_cb]
self.notify_cbs = notify_cb self.notify_cbs = notify_cb
self.disconnect_cb = disconnect_cb
async def connect(self, host, port): async def connect(self, host, port):
self.reader, self.writer = \ self.reader, self.writer = \
@ -29,6 +30,7 @@ class Receiver:
raise raise
async def close(self): async def close(self):
self.disconnect_cb = None
try: try:
self.receive_task.cancel() self.receive_task.cancel()
try: try:
@ -41,6 +43,7 @@ class Receiver:
del self.writer del self.writer
async def _receive_cr(self): async def _receive_cr(self):
try:
target = None target = None
while True: while True:
line = await self.reader.readline() line = await self.reader.readline()
@ -50,6 +53,9 @@ class Receiver:
for notify_cb in self.notify_cbs: for notify_cb in self.notify_cbs:
notify_cb(obj) notify_cb(obj)
finally:
if self.disconnect_cb is not None:
self.disconnect_cb()
class Broadcaster(AsyncioServer): class Broadcaster(AsyncioServer):

10
artiq/protocols/sync_struct.py
View File

@ -52,8 +52,10 @@ class Subscriber:
from the publisher. The mod is passed as parameter. The function is from the publisher. The mod is passed as parameter. The function is
called after the mod has been processed. called after the mod has been processed.
A list of functions may also be used, and they will be called in turn. A list of functions may also be used, and they will be called in turn.
:param disconnect_cb: An optional function called when disconnection happens
from external causes (i.e. not when ``close`` is called).
""" """
def __init__(self, notifier_name, target_builder, notify_cb=None): def __init__(self, notifier_name, target_builder, notify_cb=None, disconnect_cb=None):
self.notifier_name = notifier_name self.notifier_name = notifier_name
self.target_builder = target_builder self.target_builder = target_builder
if notify_cb is None: if notify_cb is None:
@ -61,6 +63,7 @@ class Subscriber:
if not isinstance(notify_cb, list): if not isinstance(notify_cb, list):
notify_cb = [notify_cb] notify_cb = [notify_cb]
self.notify_cbs = notify_cb self.notify_cbs = notify_cb
self.disconnect_cb = disconnect_cb
async def connect(self, host, port, before_receive_cb=None): async def connect(self, host, port, before_receive_cb=None):
self.reader, self.writer = \ self.reader, self.writer = \
@ -78,6 +81,7 @@ class Subscriber:
raise raise
async def close(self): async def close(self):
self.disconnect_cb = None
try: try:
self.receive_task.cancel() self.receive_task.cancel()
try: try:
@ -90,6 +94,7 @@ class Subscriber:
del self.writer del self.writer
async def _receive_cr(self): async def _receive_cr(self):
try:
target = None target = None
while True: while True:
line = await self.reader.readline() line = await self.reader.readline()
@ -104,6 +109,9 @@ class Subscriber:
for notify_cb in self.notify_cbs: for notify_cb in self.notify_cbs:
notify_cb(mod) notify_cb(mod)
finally:
if self.disconnect_cb is not None:
self.disconnect_cb()
class Notifier: class Notifier: