Set a known buffer size if possible on BinaryWriter()s

telethon/extensions/binary_writer.py

@@ -1,4 +1,4 @@
-from io import BufferedWriter, BytesIO
+from io import BufferedWriter, BytesIO, DEFAULT_BUFFER_SIZE
 from struct import pack
 
 
@@ -8,11 +8,16 @@ class BinaryWriter:
     Also creates a "Memory Stream" if necessary
     """
 
-    def __init__(self, stream=None):
+    def __init__(self, stream=None, known_length=None):
         if not stream:
             stream = BytesIO()
 
-        self.writer = BufferedWriter(stream)
+        if known_length is None:
+            # On some systems, DEFAULT_BUFFER_SIZE defaults to 8192
+            # That's over 16 times as big as necessary for most messages
+            known_length = max(DEFAULT_BUFFER_SIZE, 1024)
+
+        self.writer = BufferedWriter(stream, buffer_size=known_length)
         self.written_count = 0
 
     # region Writing

telethon/network/authenticator.py

@@ -19,7 +19,7 @@ def do_authentication(connection):
 
     # Step 1 sending: PQ Request
     nonce = os.urandom(16)
-    with BinaryWriter() as writer:
+    with BinaryWriter(known_length=20) as writer:
         writer.write_int(0x60469778, signed=False)  # Constructor number
         writer.write(nonce)
         sender.send(writer.get_bytes())

telethon/network/connection.py

@@ -151,8 +151,9 @@ class Connection:
     def _send_tcp_full(self, message):
         # https://core.telegram.org/mtproto#tcp-transport
         # total length, sequence number, packet and checksum (CRC32)
-        with BinaryWriter() as writer:
-            writer.write_int(len(message) + 12)
+        length = len(message) + 12
+        with BinaryWriter(known_length=length) as writer:
+            writer.write_int(length)
             writer.write_int(self._send_counter)
             writer.write(message)
             writer.write_int(crc32(writer.get_bytes()), signed=False)
@@ -160,13 +161,13 @@ class Connection:
             self.write(writer.get_bytes())
 
     def _send_intermediate(self, message):
-        with BinaryWriter() as writer:
+        with BinaryWriter(known_length=len(message) + 4) as writer:
             writer.write_int(len(message))
             writer.write(message)
             self.write(writer.get_bytes())
 
     def _send_abridged(self, message):
-        with BinaryWriter() as writer:
+        with BinaryWriter(known_length=len(message) + 4) as writer:
             length = len(message) >> 2
             if length < 127:
                 writer.write_byte(length)

telethon/network/mtproto_plain_sender.py

@@ -20,7 +20,7 @@ class MtProtoPlainSender:
 
     def send(self, data):
         """Sends a plain packet (auth_key_id = 0) containing the given message body (data)"""
-        with BinaryWriter() as writer:
+        with BinaryWriter(known_length=len(data) + 20) as writer:
             writer.write_long(0)
             writer.write_long(self._get_new_msg_id())
             writer.write_int(len(data))