Coverage for /home/agp/Documents/me/code/gutools/gutools/unet.py : 0%

Hot-keys on this page
r m x p toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1"""
2UNet module provide a asyncio Transport Layer to comunicate peers.
4This module is mainly intended to send pub/sub messages to remote peers.
6TODO: encapsulate pub/sub messages (key, payload)
8"""
9import asyncio
10import socket
11import random
12import re
13import struct
14import ujson as json
15import pickle
16import base64
17import logging
18import zlib
19import uuid
21from asyncio.protocols import DatagramProtocol
23# own libraries
24from gutools.tools import add_to_loop, new_uid, parse_uri, serializable_container, rebuild
25from gutools.uagents import UAgent, Domestic, domestic
26from gutools.uobjects import UObject, SmartContainer, Answer
29from loganalyzer.loganalyzer import setup_config
32class Request(object):
33 """Encapsulate a request with related response."""
34 def __init__(self, uri):
35 self.uri = uri
36 self.uid = new_uid()
37 self.response = Response()
39class Response(Answer):
40 """Encapsulate a response.
41 Inherits from Answer."""
42 pass
44class Session(object):
45 "TBD"
46 pass
48class UProtocolFactory():
49 """Protocol Factory used by TL."""
50 def __init__(self, uid, link, tl):
51 self.uid = uid
52 self.link = link
53 self.tl = tl
54 self.types = {
55 'broadcast': UProtocol,
56 'direct': UProtocol,
57 }
59 def __call__(self):
60 # TODO: implemente a real factory here based on config or (hash) uid
61 proto_klass = self.types.get(self.link)
62 return proto_klass(self.uid, self.tl)
65class UProtocol(DatagramProtocol, Domestic):
66 """
67 DatagramProtocol to comunicate with a peer.
69 It will keep input and output queues in order to provide a
70 reliable *stream* of datagrams.
72 User can use direct functions like `send(key, data, to, address)`
73 or publishing in the same realm with key='/net/publish'
75 - `uid`: the uid of the peer-connection.
76 - `realm`: the realm to attend pub/sub.
78 Special keys:
80 - `MISSING`: used for requesting missing datagrams
81 - `WATERMARK`: update peer watermark and drop packets received by peer.
83 (sender, to, payload)
84 """
85 def __init__(self, uid, tl):
86 DatagramProtocol.__init__(self)
87 Domestic.__init__(self)
88 self.uid = uid
89 self.tl = tl
91 self._output = dict()
92 self._input = dict()
93 self._out_sequence = dict()
94 self._in_sequence = dict()
96 self._oob_handlers = {
97 'MISSING': self._resend_missing,
98 'WATERMARK': self._update_watermark,
99 }
101 def connection_made(self, transport):
102 self.transport = transport
103 addr = self.transport.get_extra_info('socket').getsockname()
104 logging.debug("connection_made: {}: {}".format(self.uid, addr))
106 def connection_lost(self, exc):
107 logging.debug(f'connection_lost: {exc}')
109 def datagram_received(self, raw, addr):
110 baddr = self.transport.get_extra_info('socket').getsockname()
111 raw = zlib.decompress(raw)
112 sender, to, sequence, payload = json.decode(raw.decode())
113 if sender == self.uid:
114 logging.debug(f'>> [{self.uid}] Skip <{payload}> self-loop-message from {addr} in {baddr}')
115 return
117 if to in (self.uid, ): # direct message (broadcast or direct socket)
118 logging.debug(f">> [{self.uid}] Received <{payload}> from {sender}:{addr} to '{to}' in {baddr}")
120 in_sequence = self._in_sequence.get(sender, 0)
121 if sequence < in_sequence:
122 return # already dispatched message (is an echo)
124 inputs = self._input.get(sender)
125 if inputs is None:
126 inputs = self._input[sender] = dict()
127 inputs[sequence] = (payload, addr)
129 while in_sequence in inputs:
130 (key, data), addr = inputs.pop(in_sequence)
131 in_sequence += 1
132 # TODO: include sender in published message?
133 self.tl.hub.publish(key, data)
135 self._in_sequence[sender] = in_sequence
136 # self.request_missing()
138 elif to in ('', '*'): # anyone or OOB message, do not manage sequences
139 logging.debug(f">> [{self.uid}] Received <{payload}> from {sender}:{addr} to '{to}' in {baddr}")
140 (key, data) = payload
141 func = self._oob_handlers.get(key)
142 if func:
143 func(key, data, sender)
144 else:
145 self.tl.hub.publish(key, data)
146 else:
147 logging.warn(f">> [{self.uid}] Ignoring <{payload}> from {sender}:{addr} to '{to}' in {baddr}")
149 def send(self, key, data, to, address):
150 """Send a (key, data) pub/sub message to a specific peers or
151 broadcasting.
153 - If `to` is empty or '*' then message is consider broadcast.
154 - Else is added to the specific peer queue.
155 - `address` could be an <broadcast> address or direct one.
156 """
157 logging.debug(f"send: {self.uid} -> {to}: {address}")
159 # get peer sequence
160 if to not in ('', '*'):
161 sequence = self._out_sequence.setdefault(to, 0)
162 self._out_sequence[to] += 1
163 else:
164 sequence = -1
166 # encode raw message
167 payload = (key, data)
168 raw = json.encode((self.uid, to, sequence, payload)).encode()
169 raw = zlib.compress(raw)
171 # save message for replaying when network fails
172 if to not in ('', '*'):
173 output = self._output.get(to)
174 if output is None:
175 output = self._output[to] = dict()
176 output[sequence] = (raw, address)
178 self.transport.sendto(raw, address)
180 @domestic(restart=2)
181 def request_missing(self):
182 "domestic task for request missing datagrams"
183 for sender, inputs in self._input.items():
184 keys = inputs.keys()
185 if keys:
186 m = max(keys)
187 missing = set(range(self._in_sequence[sender], m))
188 missing.difference_update(keys)
189 _, address = inputs[m]
190 if missing:
191 logging.debug(f'{self.uid} Request MISSING: {missing}')
192 self.send('MISSING', missing, '', ('<broadcast>', 9999))
193 else:
194 watermark = self._in_sequence[sender] - 1
195 logging.debug(f'{self.uid} Sending WATERMARK: {watermark}')
196 self.send('WATERMARK', watermark, '', ('<broadcast>', 9999))
197 foo = 1
199 def _resend_missing(self, key, missing, sender):
200 "resent some missing datagrams"
201 logging.debug(f'{self.uid} Resend MISSING: {missing}')
202 output = self._output.get(sender)
203 if output and missing:
204 for sequence in missing:
205 info = output.get(sequence)
206 info and self.transport.sendto(*info)
207 watermark = min(missing) - 1
208 self._update_watermark(key, watermark, sender)
210 def _update_watermark(self, key, watermark, sender):
211 "Update the watermark and remove the already sent messages"
212 output = self._output.get(sender, {})
213 for k in [k for k in output.keys() if k <= watermark]:
214 output.pop(k)
216 @domestic()
217 def _shake_channels(self):
218 """Domestic function that tries to *shake* peers channels taking a
219 proactive role sending a random packet or the last sent packet
220 to the peer.
222 This behavior pretend to update remote queues in order the remote
223 peers can request for updated information.
225 This is due *missing packets* request or the *last packet* maybe can
226 be lost so peers may think there is nothing else to do, but
227 actually they have missing information that may never
228 receive unless the sender sent more packets that hopefuly update the
229 last packet sequence creating a *hole* that can be filled.
230 """
231 if random.random() > 0.5:
232 # send Last Packet (to update missing in remote site)
233 for sender, sequence in self._out_sequence.items():
234 info = self._output[sender].get(sequence - 1)
235 if info:
236 logging.debug(f'{self.uid} Sending Last: {sequence-1}')
237 self.transport.sendto(*info)
238 else:
239 # send a random packet because if more probable that
240 # remote peer already know the end of the queue
241 # so we're proactivel send a possible 'missing' packet
242 # due missing requests may also be lost
243 for sender, info in self._output.items():
244 if info:
245 sequence = random.choice(list(info.keys()))
246 info = info[sequence]
247 logging.debug(f'{self.uid} Sending Random Packet: {sequence}')
248 self.transport.sendto(*info)
252class TL(UAgent):
253 """An asyncio Transport Layer
255 This class send pub/sub message to remote peers.
257 - Message sendings are reliable.
258 - Uses two channels for UDP and UDP broadcasting.
260 User can simply publish to `/net/publish` channel in the same `realm`
261 or use `send()` method.
262 """
263 def __init__(self, uid=None, app=None, uri='>udp://:9090', loop=None):
264 super(TL, self).__init__(uid, app)
265 self.uri = uri
266 self._uri = parse_uri(uri, bind='')
267 self.channels = dict()
268 "active transports and protocols"
270 self._loop = None # set in start()
272 "Future that points out when all channels are ready for use."
274 self.app.hub.subscribe('/net/publish', self.send_pubsub)
275 self.app.hub.subscribe('/net/receive', self.receive_pubsub)
277 foo = 1
279 async def main(self):
280 """`main` asyncio method for TL agent.
282 - open channels
283 - wait until agent is stopped
284 - close channels
285 """
286 # open communications channels
287 await self.start()
288 self.channels_ready.set_result(True)
290 # now coroutine holds until Agent is stopped
291 await self.running
293 # close channels
294 await self.stop()
296 def stop(self):
297 for link, (transport, protocol, _, _) in self.channels.items():
298 protocol.stop_domestics() # del protocol ?
299 # transport.close() # done in self.stop
301 self.running.set_result(False)
304 def send_pubsub(self, key, message):
305 """Send a pub/sub (key, message) by broadcasting to all peers in
306 the same <broadcast> address.
308 TODO: need to review encapsulation. Do not use URecord.
309 """
310 # payload = (key, message)
311 raw = pickle.dumps(message)
312 raw = base64.urlsafe_b64encode(raw)
313 # raw = zlib.compress(raw)
314 transport, protocol, baddr, saddr = self.channels['broadcast']
315 protocol.send(key='/net/receive', data=raw, to='*', address=saddr)
317 def receive_pubsub(self, key, raw):
318 """Inject a remote (key, data) pub/sub message"""
319 # raw = zlib.decompress(raw)
320 raw = base64.urlsafe_b64decode(raw)
321 key, data = pickle.loads(raw)
322 self.hub.publish(key, data)
324 async def start(self):
325 """Create open channels for <broadcast> and direct communications.
327 - Create a UDP socket with specific options.
328 - Bind the socket.
329 - Create a asyncio transport and protocol.
330 - Add to self.channels dict
331 """
333 # TODO: are asyncio.events instead futures?
334 self.running = self._loop.create_future()
335 self.channels_ready = self._loop.create_future()
337 channels = {
338 'broadcast': {
339 'options': [(socket.SO_BROADCAST, 1), (socket.SO_REUSEADDR, 1), ],
340 'bind': ('', 9999),
341 },
342 'direct': {
343 'options': [(socket.SO_BROADCAST, 1), (socket.SO_REUSEADDR, 1), ],
344 }
345 }
346 address = list(self._uri['address'])
347 for i, (link, info) in enumerate(channels.items()):
348 sock = socket.socket(socket.AF_INET, socket.SOCK_DGRAM,
349 socket.SOL_UDP)
350 for option, value in info['options']:
351 sock.setsockopt(socket.SOL_SOCKET, option, value)
353 addr = info.get('bind', tuple(address))
354 sock.bind(addr)
355 # print(self.uri, link, addr)
356 # address[1] += 1 # next link will be binded to the next port
358 factory = UProtocolFactory(
359 uid=self.uid, link=link, tl=self
360 )
362 transport, protocol = await self._loop.create_datagram_endpoint(
363 lambda: factory(),
364 sock=sock,
365 )
366 if link in ('broadcast', ):
367 saddr = tuple([f'<{link}>', addr[1]])
368 else:
369 saddr = addr
370 self.channels[link] = (transport, protocol, addr, saddr)
371 foo = 1
373 async def stop(self):
374 """Close all active channels by closing its transports."""
375 for link, (transport, _, _, _) in self.channels.items():
376 transport.close()
379# HEADER_FORMAT = '>L'
380# HEADER_SIZE = struct.calcsize(HEADER_FORMAT)
381# class UStreamAgent(UAgent):
382 # def __init__(self, uid=None, realm='', uri='tcp://localhost:9090'):
383 # super().__init__(uid, realm)
384 # self.uri = uri
385 # self._uri = parse.urlparse(uri)
386 # self.address = split_netloc(self._uri.netloc)
388 # self.reader = None
389 # self.writer = None
391 # self.running = asyncio.Future()
392 # self.channels_ready = asyncio.Future()
394 # self.msg = b''
396 # async def _next_message(self, reader, writer):
397 # context = {}
398 # while len(self.msg) < HEADER_SIZE:
399 # data = await reader.read(0x1000)
400 # self.msg += data
402 # size, = struct.unpack_from(HEADER_FORMAT, self.msg)
403 # cut = size + HEADER_SIZE
404 # while len(self.msg) < cut:
405 # data = await reader.read(0x1000)
406 # self.msg += data
408 # data, self.msg = self.msg[4:cut], self.msg[cut:]
409 # data = json.decode(data)
410 # return data, context
412 # async def _send(self, data, writer):
413 # msg = bytes(json.encode(data), encoding='UTF-8')
414 # msg = struct.pack(HEADER_FORMAT, len(msg)) + msg
415 # writer.write(msg)
416 # await writer.drain()
418 # async def main(self):
419 # await self.start()
420 # self.channels_ready.set_result(True)
421 # while not self.running.done():
422 # await asyncio.sleep(1000)
423 # message, context = await self._next_message(self.reader, self.writer)
424 # self.hande(message, context)
426 # def _handle(self, message, context):
427 # print(f"receive: {message}")
428 # if message == 'quit':
429 # self.running.set_result(False)
431 # async def start(self):
432 # raise NotImplementedError('Must be overriden')
435# class UClient(UStreamAgent):
436 # """Encapsulate sending objects to UServer"""
438 # def __init__(self, uid=None, realm='', uri='tcp://localhost:9090'):
439 # super().__init__(uid, realm, uri)
441 # async def start(self):
442 # return asyncio.open_connection(*self.address)
444 # # async def main(self):
445 # # self.running = asyncio.Future()
446 # # self.reader, self.writer = await asyncio.open_connection(*self.address)
447 # # self.channels_ready.set_result(True)
448 # # while not self.running.done():
449 # # message = await self._next_message(self.reader, self.writer)
450 # # print(f"Server send: {message}")
451 # # if message == 'quit':
452 # # self.running.set_result(False)
454 # async def send(self, data):
455 # while not self.channels_ready.done():
456 # await asyncio.sleep(0.1)
457 # assert self.writer
458 # await self._send(data, self.writer)
460# class UServer(UStreamAgent):
461 # """Give access to exposed methods using TCP as well"""
463 # async def start(self):
464 # scheme = self._uri.scheme
465 # if scheme in ('tcp', ):
466 # self.server = await asyncio.start_server(self.tcp_handle, *self.address)
467 # # self.reader, self.writer = server.sockets
468 # elif scheme in ('udp', ):
469 # foo = 1
470 # else:
471 # raise RuntimeError(f'Unknown scheme {scheme}')
474 # async def main(self):
475 # self.running = asyncio.Future()
476 # async with await asyncio.start_server(self.tcp_handle, *self.address) as server:
477 # addr = server.sockets[0].getsockname()
478 # print(f'Serving on {addr}: quit to shutdown server')
479 # await self.running
480 # foo = 1
482 # async def tcp_handle(self, reader, writer):
483 # addr = writer.get_extra_info('peername')
484 # message = f"{addr!r} is connected !!!!"
485 # print(message)
486 # while not self.running.done():
487 # message = await self._next_message(reader, writer)
488 # print(f"Client send: {message}")
490 # await self._send(message, writer)
491 # if message == "exit":
492 # message = f"{addr!r} wants to close the connection."
493 # print(message)
494 # break
495 # if message == 'quit':
496 # message = f"{addr!r} wants to shutdown the server"
497 # self.running.set_result(False)
498 # writer.close()
500def test_Request():
501 uri = 'tl://121212/uagent/dbrecords/foo'
502 req = Request(uri)
504 foo = 1
507def test_parse_uri():
508 uri = parse_uri('http://myhost/mypath')
509 assert uri['scheme'] == 'http'
510 assert uri['fservice'] == 'myhost'
511 assert uri['host'] == 'myhost'
512 assert uri['path'] == '/mypath'
514 uri = parse_uri('http://agp@myhost/mypath')
515 assert uri['scheme'] == 'http'
516 assert uri['fservice'] == 'agp@myhost'
517 assert uri['host'] == 'myhost'
518 assert uri['path'] == '/mypath'
520 uri = parse_uri('http://agp:mypassword@myhost/mypath')
521 assert uri['scheme'] == 'http'
522 assert uri['fservice'] == 'agp:mypassword@myhost'
523 assert uri['host'] == 'myhost'
524 assert uri['path'] == '/mypath'
525 assert uri['user'] == 'agp'
526 assert uri['password'] == 'mypassword'
528 uri = parse_uri('http://agp:mypassword@myhost/mypath?myquery=1#42')
529 assert uri['scheme'] == 'http'
530 assert uri['fservice'] == 'agp:mypassword@myhost'
531 assert uri['host'] == 'myhost'
532 assert uri['path'] == '/mypath'
533 assert uri['user'] == 'agp'
534 assert uri['password'] == 'mypassword'
535 assert uri['query'] == 'myquery=1'
536 assert uri['fragment'] == '42'
538 uri = parse_uri('>http://agp:mypassword@myhost/mypath?myquery=1&foo=3#42')
539 assert uri['direction'] == '>'
540 assert uri['fscheme'] == '>http'
541 assert uri['scheme'] == 'http'
542 assert uri['fservice'] == 'agp:mypassword@myhost'
543 assert uri['host'] == 'myhost'
544 assert uri['path'] == '/mypath'
545 assert uri['user'] == 'agp'
546 assert uri['password'] == 'mypassword'
547 assert uri['query'] == 'myquery=1&foo=3'
548 assert uri['fragment'] == '42'
550 foo = 1
552if __name__ == '__main__':
553 test_parse_uri()
554 test_Request()