Stuff

colony/transport/framing.py

@@ -0,0 +1,32 @@
+import struct
+
+MAX_FRAME_SIZE = 16 * 1024 * 1024
+_HEADER = struct.Struct("!I")
+
+
+def _check_size(size: int):
+    if size > MAX_FRAME_SIZE:
+        raise ValueError(f"Frame of {size} bytes exceeds limit of {MAX_FRAME_SIZE}")
+
+
+def send_frame(sock, data: bytes):
+    _check_size(len(data))
+    sock.sendall(_HEADER.pack(len(data)) + data)
+
+
+def recv_frame(sock) -> bytes:
+    (size,) = _HEADER.unpack(_recv_exactly(sock, _HEADER.size))
+    _check_size(size)
+    return _recv_exactly(sock, size)
+
+
+def _recv_exactly(sock, size: int) -> bytes:
+    chunks = []
+    remaining = size
+    while remaining:
+        chunk = sock.recv(min(remaining, 65536))
+        if not chunk:
+            raise ConnectionError("Connection closed before the full message was received")
+        chunks.append(chunk)
+        remaining -= len(chunk)
+    return b"".join(chunks)

colony/utils/files.py

@@ -0,0 +1,21 @@
+import json
+import os
+import stat
+import uuid
+
+
+def atomic_write_json(path, data, indent=2):
+    directory = os.path.dirname(os.path.abspath(path))
+    tmp_path = os.path.join(directory, f".tmp-{uuid.uuid4().hex}.json")
+    fd = os.open(tmp_path, os.O_WRONLY | os.O_CREAT | os.O_EXCL, 0o666)
+    try:
+        with os.fdopen(fd, "w", encoding="utf-8") as f:
+            json.dump(data, f, indent=indent)
+            f.flush()
+            os.fsync(f.fileno())
+        if os.path.exists(path):
+            os.chmod(tmp_path, stat.S_IMODE(os.stat(path).st_mode))
+        os.replace(tmp_path, path)
+    except BaseException:
+        os.remove(tmp_path)
+        raise