reimplement bridge

salt/drone/scripts/bridge.py

 #!/usr/bin/env python3
 
-import logging
 import sys
-import re
-import datetime
-import time
+import threading
 
-import prrt
+import signal
 
-import cflib.crtp
-from cflib.crazyflie import Crazyflie, State
+from cflib.crtp.serialdriver import SerialDriver
+from cflib.crtp.prrtdriver import PrrtDriver
 from cflib.crtp.crtpstack import CRTPPacket, CRTPPort
 
-logging.basicConfig(stream=sys.stdout, level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-MTU = 32
-DEFAULT_TARGET_DELAY = 200
-PRRT_LOCAL_PORT = 5000
-
-
-class CrazyflieConnection:
-    def __init__(self, uri, receive_callback):
-        self.uri = uri
-        self.receive_callback = receive_callback
-
-        self._cf = Crazyflie()
-
-        self._cf.connected.add_callback(self._connected)
-        self._cf.disconnected.add_callback(self._disconnected)
-        self._cf.connection_failed.add_callback(self._connection_failed)
-        self._cf.connection_lost.add_callback(self._connection_lost)
-
-        logger.info('Connecting to {}'.format(uri))
-
-        # Manually open link of _cf without starting connection setup.
-        # This prevents sending messages from the library at this bridging point.
-        self._cf.connection_requested.call(uri)
-        self._cf.state = State.INITIALIZED
-        self._cf.link_uri = uri
-        try:
-            self._cf.link = cflib.crtp.get_link_driver(
-                uri, self._cf._link_quality_cb, self._cf._link_error_cb)
-
-            if not self._cf.link:
-                message = 'No driver found or malformed URI: {}'.format(uri)
-                logger.warning(message)
-                self._cf.connection_failed.call(uri, message)
-            else:
-                # Add a callback so we can check that any data is coming back from the copter
-                self._cf.packet_received.add_callback(self._cf._check_for_initial_packet_cb)
-                self._cf.platform.fetch_platform_informations(self._fetched_platform_information)
-
-        except Exception as ex:  # pylint: disable=W0703
-            # We want to catch every possible exception here and show it in the user interface
-            import traceback
-
-            logger.error("Couldn't load link driver: %s\n\n%s", ex, traceback.format_exc())
-            exception_text = "Couldn't load link driver: %s\n\n%s" % (ex, traceback.format_exc())
-            if self._cf.link:
-                self._cf.link.close()
-                self._cf.link = None
-            self._cf.connection_failed.call(uri, exception_text)
-
-        # Variable used to keep main loop occupied until disconnect
-        self.is_connected = False
-
-    def _fetched_platform_information(self):
-        self._cf.connected_ts = datetime.datetime.now()
-        self._cf.connected.call(self.uri)
-
-    def _connected(self, link_uri):
-        logger.info('Connected to {}'.format(link_uri))
-        self.is_connected = True
-        self._cf.packet_received.callbacks = [self.receive_callback]
-        self._cf.incoming.cb = []
-
-    def send(self, pk):
-        self._cf.send_packet(pk)
-
-    def _connection_failed(self, link_uri, msg):
-        """Callback when connection initial connection fails (i.e no Crazyflie at the speficied address)"""
-        logger.info('Connection to {} failed: {}'.format(link_uri, msg))
-        self.is_connected = False
-
-    def _connection_lost(self, link_uri, msg):
-        logger.info('Connection to {} lost: {}'.format(link_uri, msg))
-
-    def _disconnected(self, link_uri):
-        logger.info('Disconnected from {}'.format(link_uri))
-        self.is_connected = False
-
-    def close(self):
-        logger.info('Closing Crazyflie Connection')
-        self._cf.close_link()
-        self._cf.packet_received.remove_callback(self.receive_callback)
-        self.receive_callback = None
-
-
-class ClientConnection:
-    def __init__(self, uri):
-        self.uri = uri
-
-        uri_match = re.search(r'^prrt://((?:[\d]{1,3})\.(?:[\d]{1,3})\.(?:[\d]{1,3})\.(?:[\d]{1,3})):([\d]{1,5})'
-                              r'(?:/([\d]{1,9}))?$', uri)
-        if not uri_match:
-            raise Exception('Invalid PRRT URI')
-        address = uri_match.group(1)
-        port = int(uri_match.group(2))
-        target_delay_us = DEFAULT_TARGET_DELAY
-        if uri_match.group(3):
-            target_delay_us = int(uri_match.group(3))
-
-        logger.info('Open PRRT Link to {}:{} with target delay {}'.format(address, port, target_delay_us))
-
-        self._prrt_socket = prrt.PrrtSocket(("0.0.0.0", PRRT_LOCAL_PORT), maximum_payload_size=MTU,
-                                            target_delay=target_delay_us)
-#        self.prrt_socket.coding_configuration = prrt.PrrtCodingConfiguration(1, 1, [0])
-        self._prrt_socket.connect((address, port))
-
-    def send(self, pk):
-        pk_bytes = bytearray([pk.get_header()]) + pk.data
-        self._prrt_socket.send_sync(pk_bytes)
-
-    def receive(self):
-        pk_bytes, _ = self._prrt_socket.receive_asap()
-        if len(pk_bytes) > 0:
-            pk = CRTPPacket(pk_bytes[0], pk_bytes[1:])
-            return pk
-        else:
-            return None
-
-    def close(self):
-        logger.info('Closing Client Connection')
-        self._prrt_socket = None
-        logger.info('Disconnected from {}'.format(self.uri))
 
+class ForwardBridge(threading.Thread):
+    def __init__(self, crazyflie_connection, client_connection):
+        threading.Thread.__init__(self)
+        self._crazyflie_connection = crazyflie_connection
+        self._client_connection = client_connection
+        self.stop_running = False
 
-class Bridge:
-    def __init__(self, crazyflie_uri, client_uri):
-        cflib.crtp.init_drivers(enable_debug_driver=False)
+    def run(self):
+        while not self.stop_running:
+            pk = self._client_connection.receive_packet(-1)
+            if pk and not self.stop_running:
+                self._crazyflie_connection.send_packet(pk)
+
+    def stop(self):
+        self.stop_running = True
 
-        self._client_connection = ClientConnection(client_uri)
-        self._crazyflie_connection = CrazyflieConnection(crazyflie_uri, self._client_connection.send)
-        self.stop_running = False
 
-    def wait_for_crazyflie_connection(self):
-        while not self._crazyflie_connection.is_connected:
-            time.sleep(1)
+class BackwardBridge(threading.Thread):
+    def __init__(self, crazyflie_connection, client_connection):
+        threading.Thread.__init__(self)
+        self._crazyflie_connection = crazyflie_connection
+        self._client_connection = client_connection
+        self.stop_running = False
 
     def run(self):
-        while not self.stop_running and self._crazyflie_connection.is_connected:
-            pk = self._client_connection.receive()
-            if pk:
-                if pk.port == CRTPPort.LINKCTRL and pk.channel == 3:
-                    continue
-                self._crazyflie_connection.send(pk)
-        self.stop()
+        while not self.stop_running:
+            pk = self._crazyflie_connection.receive_packet(-1)
+            if pk and not self.stop_running:
+                self._client_connection.send_packet(pk)
 
     def stop(self):
         self.stop_running = True
+
+
+class Bridge:
+    def __init__(self, crazyflie_driver, crazyflie_uri, client_driver, client_uri):
+        print('Initializing Bridge ...')
+        self._crazyflie_connection = crazyflie_driver
+        self.crazyflie_uri = crazyflie_uri
+        self._client_connection = client_driver
+        self.client_uri = client_uri
+
+        self._crazyflie_connection.connect(crazyflie_uri, None, None)
+        self._client_connection.connect(client_uri, None, None)
+
+        self.forward_thread = None
+        self.backward_thread = None
+        print('Bridge initialized.')
+
+    def check_for_crazyflie_connection(self):
+        print('Connecting to ' + str(self.crazyflie_uri) + ' ...')
+        packet = CRTPPacket()
+        packet.port = CRTPPort.LINKCTRL
+        packet.channel = 1
+        self._crazyflie_connection.send_packet(packet)
+        self._crazyflie_connection.receive_packet(-1)
+        print('Connected to ' + str(self.crazyflie_uri) + '.')
+
+    def start(self):
+        self.forward_thread = ForwardBridge(self._crazyflie_connection, self._client_connection)
+        self.backward_thread = BackwardBridge(self._crazyflie_connection, self._client_connection)
+        self.forward_thread.daemon = True
+        self.backward_thread.daemon = True
+        self.forward_thread.start()
+        self.backward_thread.start()
+
+    def stop(self):
+        print('\nStopping threads ...')
+        self.forward_thread.stop()
+        self.backward_thread.stop()
+        print('Closing connections ...')
         self._crazyflie_connection.close()
         self._client_connection.close()
+        print('Bridge stopped.')
 
 
 if __name__ == '__main__':
+
+    def signal_handler(signal, frame):
+        bridge.stop()
+        sys.exit(0)
+
     serial_uri = 'serial://pi'
     prrt_uri = sys.argv[1]
 
-    bridge = Bridge(serial_uri, prrt_uri)
+    bridge = Bridge(SerialDriver(), serial_uri, PrrtDriver(), prrt_uri)
 
-    try:
-        bridge.wait_for_crazyflie_connection()
-        bridge.run()
-    except KeyboardInterrupt:
-        print('\n')
-        bridge.stop()
+    bridge.check_for_crazyflie_connection()
+    bridge.start()
+
+    print('\nSetup finished.\nUse "Ctrl+C" to stop.\n')
+    signal.signal(signal.SIGINT, signal_handler)
+    threading.Event().wait()