Make IPC framing more efficient (#20793)

ilevkivskyi · web-flow · commit c67a4794114d · 2026-02-12T23:20:22.000Z
This is not important now, but it will be important when we will start sending serialized ASTs over IPC, using stdlib base64 can easily add 0.5-1ms per file round-trip overhead, which is a lot. Initially I waned to use `librt.base64`, but ultimately simply switched to framing with a fixed size header. I am still using `librt.base64` in few places where performance doesn't really matter, because why not. Btw I noticed we have very few mypyc primitives for `bytearray`, probably we should add more, added comment in mypyc/mypyc#644
diff --git a/mypy-requirements.txt b/mypy-requirements.txt
@@ -4,4 +4,4 @@ typing_extensions>=4.6.0
 mypy_extensions>=1.0.0
 pathspec>=1.0.0
 tomli>=1.1.0; python_version<'3.11'
-librt>=0.6.2; platform_python_implementation != 'PyPy'
+librt>=0.8.0; platform_python_implementation != 'PyPy'
diff --git a/mypy/build.py b/mypy/build.py
@@ -13,7 +13,6 @@
 
 from __future__ import annotations
 
-import base64
 import collections
 import contextlib
 import gc
@@ -40,6 +39,7 @@
     TypedDict,
 )
 
+from librt.base64 import b64encode
 from librt.internal import (
     cache_version,
     read_bool,
@@ -353,7 +353,7 @@ def default_flush_errors(
     if options.num_workers > 0:
         # TODO: switch to something more efficient than pickle (also in the daemon).
         pickled_options = pickle.dumps(options.snapshot())
-        options_data = base64.b64encode(pickled_options).decode()
+        options_data = b64encode(pickled_options).decode()
         workers = [
             WorkerClient(f".mypy_worker.{idx}.json", options_data, worker_env or os.environ)
             for idx in range(options.num_workers)
diff --git a/mypy/build_worker/worker.py b/mypy/build_worker/worker.py
@@ -14,7 +14,6 @@
 from __future__ import annotations
 
 import argparse
-import base64
 import gc
 import json
 import os
@@ -24,6 +23,8 @@
 import time
 from typing import NamedTuple
 
+from librt.base64 import b64decode
+
 from mypy import util
 from mypy.build import (
     AckMessage,
@@ -72,7 +73,7 @@ def main(argv: list[str]) -> None:
     # This mimics how daemon receives the options. Note we need to postpone
     # processing error codes after plugins are loaded, because plugins can add
     # custom error codes.
-    options_dict = pickle.loads(base64.b64decode(args.options_data))
+    options_dict = pickle.loads(b64decode(args.options_data))
     options_obj = Options()
     disable_error_code = options_dict.pop("disable_error_code", [])
     enable_error_code = options_dict.pop("enable_error_code", [])
diff --git a/mypy/dmypy/client.py b/mypy/dmypy/client.py
@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import argparse
-import base64
 import json
 import os
 import pickle
@@ -17,6 +16,8 @@
 from collections.abc import Callable, Mapping
 from typing import Any, NoReturn
 
+from librt.base64 import b64decode
+
 from mypy.defaults import RECURSION_LIMIT
 from mypy.dmypy_os import alive, kill
 from mypy.dmypy_util import DEFAULT_STATUS_FILE, receive, send
@@ -620,7 +621,7 @@ def do_daemon(args: argparse.Namespace) -> None:
     if args.options_data:
         from mypy.options import Options
 
-        options_dict = pickle.loads(base64.b64decode(args.options_data))
+        options_dict = pickle.loads(b64decode(args.options_data))
         options_obj = Options()
         options = options_obj.apply_changes(options_dict)
     else:
diff --git a/mypy/dmypy_server.py b/mypy/dmypy_server.py
@@ -7,7 +7,6 @@
 from __future__ import annotations
 
 import argparse
-import base64
 import io
 import json
 import os
@@ -20,6 +19,8 @@
 from contextlib import redirect_stderr, redirect_stdout
 from typing import Any, Final, TypeAlias as _TypeAlias
 
+from librt.base64 import b64encode
+
 import mypy.build
 import mypy.errors
 import mypy.main
@@ -57,7 +58,7 @@ def daemonize(
         """
         command = [sys.executable, "-m", "mypy.dmypy", "--status-file", status_file, "daemon"]
         pickled_options = pickle.dumps(options.snapshot())
-        command.append(f'--options-data="{base64.b64encode(pickled_options).decode()}"')
+        command.append(f'--options-data="{b64encode(pickled_options).decode()}"')
         if timeout:
             command.append(f"--timeout={timeout}")
         if log_file:
diff --git a/mypy/ipc.py b/mypy/ipc.py
@@ -1,16 +1,15 @@
-"""Cross platform abstractions for inter-process communication
+"""Cross-platform abstractions for inter-process communication
 
 On Unix, this uses AF_UNIX sockets.
 On Windows, this uses NamedPipes.
 """
 
 from __future__ import annotations
 
-import base64
-import codecs
 import json
 import os
 import shutil
+import struct
 import sys
 import tempfile
 from abc import abstractmethod
@@ -20,6 +19,7 @@
 from typing import Final
 from typing_extensions import Self
 
+from librt.base64 import urlsafe_b64encode
 from librt.internal import ReadBuffer, WriteBuffer
 
 if sys.platform == "win32":
@@ -37,6 +37,9 @@
 
     _IPCHandle = socket.socket
 
+# Size of the message packed as !L, i.e. 4 bytes in network order (big-endian).
+HEADER_SIZE = 4
+
 
 class IPCException(Exception):
     """Exception for IPC issues."""
@@ -58,24 +61,29 @@ class IPCBase:
     def __init__(self, name: str, timeout: float | None) -> None:
         self.name = name
         self.timeout = timeout
+        self.message_size: int | None = None
         self.buffer = bytearray()
 
-    def frame_from_buffer(self) -> bytearray | None:
+    def frame_from_buffer(self) -> bytes | None:
         """Return a full frame from the bytes we have in the buffer."""
-        space_pos = self.buffer.find(b" ")
-        if space_pos == -1:
+        size = len(self.buffer)
+        if size < HEADER_SIZE:
             return None
-        # We have a full frame
-        bdata = self.buffer[:space_pos]
-        self.buffer = self.buffer[space_pos + 1 :]
-        return bdata
+        if self.message_size is None:
+            self.message_size = struct.unpack("!L", self.buffer[:HEADER_SIZE])[0]
+        if size < self.message_size + HEADER_SIZE:
+            return None
+        # We have a full frame, avoid extra copy in case we get a large frame.
+        bdata = memoryview(self.buffer)[HEADER_SIZE : HEADER_SIZE + self.message_size]
+        self.buffer = self.buffer[HEADER_SIZE + self.message_size :]
+        self.message_size = None
+        return bytes(bdata)
 
     def read(self, size: int = 100000) -> str:
         return self.read_bytes(size).decode("utf-8")
 
     def read_bytes(self, size: int = 100000) -> bytes:
         """Read bytes from an IPC connection until we have a full frame."""
-        bdata: bytearray | None = bytearray()
         if sys.platform == "win32":
             while True:
                 # Check if we already have a message in the buffer before
@@ -126,19 +134,19 @@ def read_bytes(self, size: int = 100000) -> bytes:
                 self.buffer.extend(more)
 
         if not bdata:
-            # Socket was empty and we didn't get any frame.
+            # Socket was empty, and we didn't get any frame.
             # This should only happen if the socket was closed.
             return b""
-        return codecs.decode(bdata, "base64")
+        return bdata
 
     def write(self, data: str) -> None:
         self.write_bytes(data.encode("utf-8"))
 
     def write_bytes(self, data: bytes) -> None:
         """Write to an IPC connection."""
 
-        # Frame the data by urlencoding it and separating by space.
-        encoded_data = codecs.encode(data, "base64") + b" "
+        # Frame the data by adding fixed size header.
+        encoded_data = struct.pack("!L", len(data)) + data
 
         if sys.platform == "win32":
             try:
@@ -226,9 +234,7 @@ class IPCServer(IPCBase):
 
     def __init__(self, name: str, timeout: float | None = None) -> None:
         if sys.platform == "win32":
-            name = r"\\.\pipe\{}-{}.pipe".format(
-                name, base64.urlsafe_b64encode(os.urandom(6)).decode()
-            )
+            name = r"\\.\pipe\{}-{}.pipe".format(name, urlsafe_b64encode(os.urandom(6)).decode())
         else:
             name = f"{name}.sock"
         super().__init__(name, timeout)
diff --git a/pyproject.toml b/pyproject.toml
@@ -9,7 +9,7 @@ requires = [
     "mypy_extensions>=1.0.0",
     "pathspec>=1.0.0",
     "tomli>=1.1.0; python_version<'3.11'",
-    "librt>=0.6.2; platform_python_implementation != 'PyPy'",
+    "librt>=0.8.0; platform_python_implementation != 'PyPy'",
     # the following is from build-requirements.txt
     "types-psutil",
     "types-setuptools",
@@ -53,7 +53,7 @@ dependencies = [
   "mypy_extensions>=1.0.0",
   "pathspec>=1.0.0",
   "tomli>=1.1.0; python_version<'3.11'",
-  "librt>=0.6.2; platform_python_implementation != 'PyPy'",
+  "librt>=0.8.0; platform_python_implementation != 'PyPy'",
 ]
 dynamic = ["version"]
 
diff --git a/test-requirements.txt b/test-requirements.txt
@@ -22,7 +22,7 @@ identify==2.6.15
     # via pre-commit
 iniconfig==2.1.0
     # via pytest
-librt==0.7.8 ; platform_python_implementation != "PyPy"
+librt==0.8.0 ; platform_python_implementation != "PyPy"
     # via -r mypy-requirements.txt
 lxml==6.0.2 ; python_version < "3.15"
     # via -r test-requirements.in
