Changeset a574a93 in ammosreader

Timestamp:

05/09/22 06:15:46 (3 years ago)

Author:

recknagel <recknagel@…>

Branches:

AmmosSource, guix

Children:

b67e7e5

Parents:

4180d6a

Message:

added error handling for timed out connections

File:

• src/ammosreader/AmmosAudioSocketReader.py (modified) (6 diffs)

src/ammosreader/AmmosAudioSocketReader.py

@@ -1 +1 @@
 """I read a Ammos datastream from a socket."""
 
+import select
 import socket
 from collections import deque
@@ -10 +11 @@
 
 class AmmosAudioSocketReader:
-    def __init__(self, socket:socket.socket):
+    def __init__(self, socket:socket.socket, debug=True):
         """
         Initializes the AmmosAudioSocketReader
@@ -16 +17 @@
         Args:
             socket (socket.socket): socket to read from
+            debug (bool): if true, prints debug information
         """
 
@@ -24 +26 @@
         self.__socket = socket
 
-    def __read_next_audio_data_body(self, sample_count:int, channel_count:int, sample_size:int) -> bytearray:
+        #
+        self.DEBUG_MODE = debug
+
+    def __get_next_data(self, byte_count: int) -> bytearray:
         """
-        reads the next audio data body
+        Gets the next bytes from the socket, for example headers and body data.
 
         Args:
-            sample_count (int): amount of samples
-            channel_count (int): amount of channels
-            sample_size (int): size of a sample in bytes
+            byte_count (int): number of bytes to read
+
+        Raises:
+            TimeoutError: Raises TimeoutError if the socket does not serve data anymore
 
         Returns:
-            bytearray: contains the audio data
+            bytearray: data from socket as bytearray
         """
 
-        total = sample_count*channel_count*sample_size
         byte_array = []
+        
+        while len(b''.join(byte_array)) < byte_count:
+            if self.DEBUG_MODE:
+                print(f"Remaining Bytes: {byte_count - len(b''.join(byte_array))}")
+            self.__socket.settimeout(5)
+            new_bytes = self.__socket.recv(byte_count - len(b''.join(byte_array)), socket.MSG_WAITALL)
 
-        while len(b''.join(byte_array)) < total:
-            byte_array.append(self.__socket.recv(total - len(b''.join(byte_array))))
+            if not new_bytes:
+                raise TimeoutError("Socket timed out while reading data")
 
-        if len(b''.join(byte_array)) != total:
-            print("Can not read all", total, "bytes of data body")
-            return None
+            if self.DEBUG_MODE:
+                print(f"Got {len(new_bytes)} bytes of {byte_count - len(b''.join(byte_array))} ramining")
+
+            byte_array.append(new_bytes)
+
         return b''.join(byte_array)
 
@@ -61 +74 @@
         return np.frombuffer(audio_data_body, dtype=np.int16)
 
-    def read_next_frame(self) -> tuple[bytearray, int]:
-        """
-        reads the next ammos audio frame
+    def read_next_frame(self) -> tuple[np.ndarray, int]:
+        """Reads the next ammos audio frame from socket
+
+        Raises:
+            TimeoutError: Raisees TimeoutError if the socket does not serve data anymore
 
         Returns:
-            tuple[bytearray, int]: contains the audio data and the sample rate
+            tuple[np.ndarray, int]: Contains the audio data and the sample rate
         """
+
+        # get first byte of the day
+        self.__socket.settimeout(5)
+
+        new_byte = self.__socket.recv(1, socket.MSG_WAITALL)
+        # raise Exception if socket does not return anything
+        if len(new_byte) < 1:
+            raise TimeoutError      
+
         #read loop
-        byte = self.__socket.recv(1)
-
-        while byte:
+        while new_byte:
             #
-            self.__magic_word_buffer.append(byte)
+            self.__magic_word_buffer.append(new_byte)
             byte_array = b''.join(self.__magic_word_buffer)
 
@@ -80 +102 @@
 
                 ammos_global_header_buffer = list(self.__magic_word_buffer)
-                while len(b''.join(ammos_global_header_buffer)) < 24:
-                    ammos_global_header_buffer.append(self.__socket.recv(24 - len(b''.join(ammos_global_header_buffer))))
+                ammos_global_header_buffer.append(self.__get_next_data(20))
+                #while len(b''.join(ammos_global_header_buffer)) < 24:
+                #    ammos_global_header_buffer.append(self.__socket.recv(24 - len(b''.join(ammos_global_header_buffer))))
                     
                 ammos_global_header = AmmosGlobalFrameHeader.from_bytes(b''.join(ammos_global_header_buffer))
-                print(ammos_global_header)
+                if self.DEBUG_MODE:
+                    print(ammos_global_header)
 
                 if ammos_global_header.data_header_length == 44 and ammos_global_header.frame_type == 256:
-                    byte_array_header = []
-                    while len(b''.join(byte_array_header)) < 44:
-                        byte_array_header.append(self.__socket.recv(44 - len(b''.join(byte_array_header))))
+                    byte_array_header = self.__get_next_data(44)
+                    #while len(b''.join(byte_array_header)) < 44:
+                    #    byte_array_header.append(self.__socket.recv(44 - len(b''.join(byte_array_header))))
 
                     ammos_extended_audio_data_header = AmmosExtendedAudioDataHeader.from_bytes(b''.join(byte_array_header))
-                    print(ammos_extended_audio_data_header.sample_count, ammos_extended_audio_data_header.channel_count, ammos_extended_audio_data_header.sample_size)
-                    audio_body = self.__read_next_audio_data_body(ammos_extended_audio_data_header.sample_count, 
-                                                                  ammos_extended_audio_data_header.channel_count, 
-                                                                  ammos_extended_audio_data_header.sample_size)
+                    if self.DEBUG_MODE:
+                        print(ammos_extended_audio_data_header.sample_count, ammos_extended_audio_data_header.channel_count, ammos_extended_audio_data_header.sample_size)
+                    audio_body = self.__get_next_data(ammos_extended_audio_data_header.sample_count* 
+                                                      ammos_extended_audio_data_header.channel_count* 
+                                                      ammos_extended_audio_data_header.sample_size)
 
                     audio_array = self.__audio_data_body_to_numpy(audio_body)
-                    print(len(audio_array), len(audio_array)/ammos_extended_audio_data_header.sample_rate)
+                    if self.DEBUG_MODE:
+                        print(len(audio_array), len(audio_array)/ammos_extended_audio_data_header.sample_rate)
 
                     return [audio_array, ammos_extended_audio_data_header.sample_rate]
 
                 elif ammos_global_header.data_header_length == 36 and ammos_global_header.frame_type == 256:
-                    byte_array_header = []
-                    while len(b''.join(byte_array_header)) < 36:
-                        byte_array_header.append(self.__socket.recv(36 - len(b''.join(byte_array_header))))
+                    byte_array_header = self.__get_next_data(36)
+                    #while len(b''.join(byte_array_header)) < 36:
+                    #    byte_array_header.append(self.__socket.recv(36 - len(b''.join(byte_array_header))))
 
                     ammos_audio_data_header = AmmosAudioDataHeader.from_bytes(b''.join(byte_array_header))
-                    print(ammos_audio_data_header.sample_count, ammos_audio_data_header.channel_count, ammos_audio_data_header.sample_size)
-                    audio_body = self.__read_next_audio_data_body(ammos_audio_data_header.sample_count, 
-                                                                  ammos_audio_data_header.channel_count, 
-                                                                  ammos_audio_data_header.sample_size)
+                    if self.DEBUG_MODE:
+                        print(ammos_audio_data_header.sample_count, ammos_audio_data_header.channel_count, ammos_audio_data_header.sample_size)
+                    audio_body = self.__get_next_data(ammos_extended_audio_data_header.sample_count* 
+                                                      ammos_extended_audio_data_header.channel_count* 
+                                                      ammos_extended_audio_data_header.sample_size)
 
                     audio_array = self.__audio_data_body_to_numpy(audio_body)
-                    print(len(audio_array), len(audio_array)/ammos_audio_data_header.sample_rate)
+                    if self.DEBUG_MODE:
+                        print(len(audio_array), len(audio_array)/ammos_audio_data_header.sample_rate)
 
                     return [audio_array, ammos_audio_data_header.sample_rate]
 
-            byte = self.__socket.recv(1)
+            # get the next byte
+            self.__socket.settimeout(5)
+
+            new_byte = self.__socket.recv(1, socket.MSG_WAITALL)
+            # raise Exception if socket does not return anything
+            if len(new_byte) < 1:
+                raise TimeoutError   
 
         return None