def setup_nodes(self, count=2):
_logger.debug('DHTNetwork setup---')
self.er = []
self.er_thread = []
self.nc = []
self.timer = []
for i in range(count):
self.er.append(EventReactor())
er_thread = threading.Thread(target=self.er[i].start)
er_thread.daemon = True
er_thread.name = 'event-reactor-thread-%d' % i
er_thread.start()
self.er_thread.append(er_thread)
self.nc.append(DHTNetwork(self.er[i], MemoryKVPTable()))
timer = EventScheduler(self.er[i])
timer.add_one_shot(self.TIMEOUT, EventReactor.STOP_ID)
self.timer.append(timer)
_logger.debug('Server %d at %s', i, self.nc[i].address)
self.stuff = {}
class Network(EventReactorMixin):
'''Network controller
:CVariables:
MAX_UDP_PACKET_SIZE
The maximum UDP packet size allowed
DEFAULT_TIMEOUT
The time in seconds before a reply is timed out
STREAM_DATA_SIZE
The size in bytes of the parts of the file transmitted
'''
MAX_UDP_PACKET_SIZE = 65507 # bytes
DEFAULT_TIMEOUT = 10 # seconds
STREAM_DATA_SIZE = 1024 # bytes
SEQUENCE_ID_SIZE = 20 # bytes
DEFAULT_POOL_SIZE = 20
def __init__(self, event_reactor, address=('127.0.0.1', 0)):
EventReactorMixin.__init__(self, event_reactor)
self._server = UDPServer(event_reactor, address=address)
# By passing in the same socket object to the client, this method
# allows other nodes to reply to our server's port.
self._client = UDPClient(socket_obj=self._server.socket)
self._reply_table = ReplyTable()
self._downloads = {}
self._pool_executor = WrappedThreadPoolExecutor(
Network.DEFAULT_POOL_SIZE, event_reactor)
self._event_scheduler = EventScheduler(event_reactor)
self._transfer_timer_id = EventID(self, 'Clean transfers')
self._running = True
self._register_handlers()
self._server.start()
@property
def server_address(self):
'''The address of the server'''
return self._server.server_address
def _register_handlers(self):
'''Register the event callbacks'''
self.event_reactor.register_handler(UDP_INBOUND_EVENT,
self._udp_incoming_callback)
self.event_reactor.register_handler(EventReactor.STOP_ID,
self._stop_callback)
self.event_reactor.register_handler(self._transfer_timer_id,
self._clean_download)
def _stop_callback(self, event_id):
'''Stop and expire everything'''
self._running = False
for transfer_id in list(self._downloads.keys()):
download_task = self._downloads[transfer_id]
del self._downloads[transfer_id]
download_task.transfer(None)
for key in list(self._reply_table.out_table.keys()):
event = self._reply_table.out_table[key]
event.set()
def _clean_download(self, event_id, transfer_id):
'''Remove timed out file download'''
download_task = self._downloads[transfer_id]
last_modified = download_task.last_modified
timeout = download_task.timeout
if last_modified + timeout < time.time():
_logger.debug('Cleaned out download %s', transfer_id)
del self._downloads[transfer_id]
download_task.transfer(None)
else:
_logger.debug('Still alive download %s', transfer_id)
self._event_scheduler.add_one_shot(timeout,
self._transfer_timer_id, transfer_id)
def _udp_incoming_callback(self, event_id, address, data):
'''udp incoming'''
if not self._running:
return
_logger.debug('UDP %s←%s %s', self.server_address, address,
data[:160])
packet_dict = self._unpack_udp_data(data)
if not packet_dict:
return
data_packet = DataPacket(address, packet_dict,
packet_dict.get(JSONKeys.SEQUENCE_ID) \
or packet_dict.get(JSONKeys.REPLY_SEQUENCE_ID))
if JSONKeys.REPLY_SEQUENCE_ID in packet_dict:
self._accept_reply(data_packet)
elif JSONKeys.TRANSFER_ID in packet_dict:
self._accept_transfer(data_packet)
else:
self._accept_packet(data_packet)
def _accept_packet(self, data_packet):
self.receive_callback(data_packet)
def receive_callback(self, data_packet):
'''The function called when a data packet arrives.
:Parameters:
data_packet: :class:`DataPacket`
The incoming data packet
This function is called for packets that are not replies. Implementors
of this class should override this method.
'''
raise NotImplementedError()
def expect_incoming_transfer(self, transfer_id, timeout=DEFAULT_TIMEOUT,
download_task_class=None, max_size=None):
'''Allow a transfer for download.
:Parameters:
transfer_id: ``str``
A transfer id that the other client use for transferring data.
timeout: ``int`` ``float``
Time in seconds before the transfer times out.
max_size: ``int`` ``None``
The maximum file size.
:rtype: :class:`DownloadTask`
:return: A future that returns a file object that may have been
interrupted. The progress is the number of bytes downloaded.
'''
download_task_class = download_task_class or DownloadTask
download_task = download_task_class(max_size=max_size)
self._downloads[transfer_id] = download_task
self._event_scheduler.add_one_shot(timeout, self._transfer_timer_id,
transfer_id)
self._pool_executor.submit(download_task)
return download_task
def _accept_reply(self, data_packet):
'''Process a reply and allow a future to resume'''
sequence_id = data_packet.sequence_id
address = data_packet.address
event = self._reply_table.get_out_entry(sequence_id,
address)
if not event:
_logger.debug('Unknown seq id %s, packet discarded', sequence_id)
return
self._reply_table.remove_out_entry(sequence_id, address)
self._reply_table.add_in_entry(sequence_id, address, data_packet)
event.set()
def _accept_transfer(self, data_packet):
'''Process a file download'''
transfer_id = data_packet.dict_obj[JSONKeys.TRANSFER_ID]
if transfer_id in self._downloads:
if JSONKeys.TRANSFER_DATA in data_packet.dict_obj:
self._read_download(data_packet, transfer_id)
return
_logger.debug('Transfer discarded')
def _read_download(self, data_packet, transfer_id):
'''Read data'''
_logger.debug('Read download')
download_task = self._downloads[transfer_id]
data_str = data_packet.dict_obj[JSONKeys.TRANSFER_DATA]
download_task.address = data_packet.address
if data_str is None:
download_task.transfer(None)
_logger.debug('Read download finished')
return
else:
try:
data = base64.b64decode(data_str.encode())
except binascii.Error as e:
_logger.debug('Decode error %s', e)
return
download_task.transfer(data)
_logger.debug('Read download len=%d', len(data))
if download_task.is_running:
d = {
JSONKeys.TRANSFER_ID: transfer_id
}
self.send_answer_reply(data_packet, d)
else:
_logger.debug('Download aborted')
def _pack_udp_data(self, packet_dict):
'''Pack the dict into a format suitable for transmission.
The format currently is JSON.
'''
data = zlib.compress(json.dumps(packet_dict).encode())
if len(data) < Network.MAX_UDP_PACKET_SIZE:
_logger.debug('Packed data %s', data[:20])
return data
else:
raise Exception('data size too large')
def _unpack_udp_data(self, data):
'''Convert the data into a dict'''
try:
dict_obj = json.loads(zlib.decompress(data).decode())
except Exception as e:
_logger.debug('Failed json parsing %s', e)
return
if not isinstance(dict_obj, dict):
_logger.debug('Not a dict')
return
return dict_obj
def send(self, address, dict_obj, timeout=None):
'''Send the ``dict`` to address
:Parameters:
address: ``tuple``
A 2-tuple with the host and port number.
dict_obj: ``dict``
The ``dict`` that will be converted to JSON format.
timeout: ``None``, ``int``, ``float``, ``True``
If `timeout` is a number, the class will attempt to ensure
delivery and wait for a reply. A future will be returned.
If ``True``, the default timeout will be used.
:rtype: ``None``, :class:`SendPacketTask`
:return: Returns a :class:`SendPacketTask` if timeout is given.
The result is either :class:`DataPacket` or ``None``.
'''
if timeout is None:
self._send_plain(address, dict_obj)
else:
if timeout is True:
timeout = Network.DEFAULT_TIMEOUT
return self._send_expect_reply(address, dict_obj, timeout)
def _send_plain(self, address, dict_obj):
'''Send the data as a single UDP packet'''
_logger.debug('Dict %s→%s', self.server_address, address)
self._client.send(address, self._pack_udp_data(dict_obj))
def _send_expect_reply(self, address, dict_obj, timeout=DEFAULT_TIMEOUT):
'''Send the data and wait for a reply
:rtype: :class:`SendPacketTask`
'''
_logger.debug('Dict %s→%s timeout=%d', self.server_address,
address, timeout)
sequence_id = self.new_sequence_id()
event = threading.Event()
self._reply_table.add_out_entry(sequence_id, address, event)
packet_dict = dict_obj.copy()
packet_dict[JSONKeys.SEQUENCE_ID] = sequence_id
def send_fn():
self._client.send(address, self._pack_udp_data(packet_dict))
send_packet_task = SendPacketTask(send_fn, sequence_id, address,
self._reply_table, event, timeout)
self._pool_executor.submit(send_packet_task)
return send_packet_task
def send_answer_reply(self, source_data_packet, dict_obj):
'''Send ``dict`` that is a response to a incoming data packet
:Parameters:
source_data_packet: :class:`DataPacket`
The original incoming data packet to respond to.
dict_obj: ``dict``
The data to send back
Use this function to reply to packets that expect a response. This
function automatically adds sequence IDs the reply packet.
'''
address = source_data_packet.address
sequence_id = source_data_packet.sequence_id
_logger.debug('Dict reply %s→%s seq_id=%s', self.server_address,
address, sequence_id)
packet_dict = dict_obj.copy()
packet_dict[JSONKeys.REPLY_SEQUENCE_ID] = sequence_id
self._client.send(address, self._pack_udp_data(packet_dict))
def send_bytes(self, address, transfer_id, bytes_,
timeout=DEFAULT_TIMEOUT):
'''Transfer data to another client.
:Parameters:
address: ``tuple``
A 2-tuple with host and port number.
bytes_: ``bytes``
The data to send.
timeout: ``int``, ``float``
The time in seconds before the transfer times out.
transfer_id: ``str``, ``None``
The transfer ID to be used. If ``None``, an ID will be
created automatically.
:see: :func:`send_file`
:rtype: :class:`UploadTask`
'''
f = io.BytesIO(bytes_)
return self.send_file(address, transfer_id, f, timeout)
def send_file(self, address, transfer_id, file_, timeout=DEFAULT_TIMEOUT):
'''Transfer data to another client.
:Parameters:
address: ``tuple``
A 2-tuple with host and port number.
file_: ``str``, ``object``
A filename or a file-like object which has ``read``.
timeout: ``int``, ``float``
The time in seconds before the transfer times out.
transfer_id: ``str``, ``None``
The transfer ID to be used. If ``None``, an ID will be
created automatically.
:rtype: :class:`UploadTask`
:return: A future that returns an ``int`` that is the number of bytes
sent.
'''
if hasattr(file_, 'read'):
source_file = file_
else:
source_file = open(file_, 'rb')
transfer_id = transfer_id or self.new_sequence_id()
_logger.debug('Send file %s→%s', self.server_address, address)
upload_task = UploadTask(self, address, source_file,
transfer_id, timeout)
self._pool_executor.submit(upload_task)
return upload_task
def new_sequence_id(self):
'''Generate a new sequence ID.
:rtype: ``str``
'''
return bytes_to_b64(os.urandom(Network.SEQUENCE_ID_SIZE))
