Optimizing UDP Transmission for High-Latency Networks in Python

Question

I need to create a Python UDP Utility for Linux, but I can't optimize for packet loss simulations. It working great and actually sending the file, but its working really slow.

This is my code:

client-side

import socket
import struct
import time
import sys

# Constants
SERVER_ADDRESS = ('localhost', 12345)
CHUNK_SIZE = 1024
HEADER_FORMAT = "I"
INITIAL_TIMEOUT = 0.5
WINDOW_SIZE = 4
MIN_TIMEOUT = 0.1
MAX_TIMEOUT = 2.0

def calculate_new_timeout(estimated_rtt, sample_rtt, alpha=0.125, beta=0.25):
    if estimated_rtt is None:
        estimated_rtt = sample_rtt
    else:
        estimated_rtt = (1 - alpha) * estimated_rtt + alpha * sample_rtt
    adjusted_timeout = estimated_rtt + beta * estimated_rtt
    return max(MIN_TIMEOUT, min(adjusted_timeout, MAX_TIMEOUT))

def send_chunk(s, chunk, seq_num):
    try:
        packet = struct.pack(HEADER_FORMAT, seq_num) + chunk
        s.sendto(packet, SERVER_ADDRESS)
    except socket.error as err:
        print(f"Error sending packet {seq_num}: {err}")

def send_file(filename):
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.settimeout(INITIAL_TIMEOUT)
    estimated_rtt = None
    packets_in_flight = {}

    start_time = time.time()

    try:
        with open(filename, 'rb') as f:
            eof = False
            seq_num = 0
            while not eof:
                while len(packets_in_flight) < WINDOW_SIZE and not eof:
                    chunk = f.read(CHUNK_SIZE)
                    if chunk == b'':
                        eof = True
                        chunk = b'EOF'
                    send_chunk(s, chunk, seq_num)
                    packets_in_flight[seq_num] = (chunk, time.time())
                    seq_num += 1

                while packets_in_flight:
                    try:
                        ack, _ = s.recvfrom(CHUNK_SIZE)
                        ack_num, = struct.unpack(HEADER_FORMAT, ack)
                        if ack_num in packets_in_flight:
                            sample_rtt = time.time() - packets_in_flight[ack_num][1]
                            del packets_in_flight[ack_num]
                            estimated_rtt = calculate_new_timeout(estimated_rtt, sample_rtt)
                            s.settimeout(estimated_rtt)
                    except socket.timeout:
                        for seq_num, (chunk, _) in packets_in_flight.items():
                            send_chunk(s, chunk, seq_num)
    finally:
        duration = time.time() - start_time
        print(f'File sent successfully in {duration:.2f} seconds.')
        s.close()

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python client.py <filename>")
        sys.exit(1)
    send_file(sys.argv[1])

Sever-side

import socket
import struct
import sys

# Constants
LISTEN_ADDRESS = ('localhost', 12345)
BUFFER_SIZE = 2048
HEADER_FORMAT = "I"

def receive_file(save_as):
    s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
    s.bind(LISTEN_ADDRESS)
    print("Server started. Waiting for data...")

    expected_seq_num = 0

    try:
        with open(save_as, 'wb') as f:
            while True:
                packet, client_address = s.recvfrom(BUFFER_SIZE)
                seq_num, = struct.unpack(HEADER_FORMAT, packet[:struct.calcsize(HEADER_FORMAT)])
                data = packet[struct.calcsize(HEADER_FORMAT):]

                if seq_num == expected_seq_num:
                    if data == b'EOF':
                        print("EOF received. Ending transmission.")
                        ack_packet = struct.pack(HEADER_FORMAT, expected_seq_num)
                        s.sendto(ack_packet, client_address)
                        break
                    f.write(data)
                    expected_seq_num += 1

                ack_packet = struct.pack(HEADER_FORMAT, expected_seq_num - 1)
                s.sendto(ack_packet, client_address)

    finally:
        s.close()
        print(f"File received successfully and saved as {save_as}.")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python server.py <save_as_filename>")
        sys.exit(1)
    receive_file(sys.argv[1])

I already tried implement Go-Back-N algorithm and dynamic timeout, but still its not working fast enough. Maybe anyone knows another way to optimize it?