From e88cbb016281c79088a6a56b32ba737b72a4c045 Mon Sep 17 00:00:00 2001
From: Shane Smiskol <shane@smiskol.com>
Date: Wed, 12 Jun 2024 21:14:23 -0700
Subject: [PATCH] add zstd Python library (#32731)

* add zstd

* fix

* uploader

* logreader: zst support

* caps

* only logreader
old-commit-hash: 04a93dd849461df7a3496706c6172d95e83149c9
---
 poetry.lock            | 4 ++--
 pyproject.toml         | 1 +
 tools/lib/logreader.py | 8 ++++++--
 3 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/poetry.lock b/poetry.lock
index 06d2ff3162..532a43c6a1 100644
--- a/poetry.lock
+++ b/poetry.lock
@@ -1,3 +1,3 @@
 version https://git-lfs.github.com/spec/v1
-oid sha256:6181d338c8da4314745ebb305485a8c22130288f3faa9740fb1782288c06bc0b
-size 614543
+oid sha256:35ddb7f95b0443e48f540d6b7ec0b01b4d35da2c7b758de47fb6f0040dac527d
+size 627912
diff --git a/pyproject.toml b/pyproject.toml
index 7b780c87f0..edc515771a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -187,6 +187,7 @@ tabulate = "*"
 types-requests = "*"
 types-tabulate = "*"
 tqdm = "*"
+zstd = "*"
 
 # this is only pinned since 5.15.11 is broken
 pyqt5 = { version = "==5.15.2", markers = "platform_machine == 'x86_64'" } # no aarch64 wheels for macOS/linux
diff --git a/tools/lib/logreader.py b/tools/lib/logreader.py
index 48a04b28de..2430f1542c 100755
--- a/tools/lib/logreader.py
+++ b/tools/lib/logreader.py
@@ -10,6 +10,7 @@ import sys
 import tqdm
 import urllib.parse
 import warnings
+import zstd
 
 from collections.abc import Callable, Iterable, Iterator
 from urllib.parse import parse_qs, urlparse
@@ -34,8 +35,8 @@ class _LogFileReader:
     ext = None
     if not dat:
       _, ext = os.path.splitext(urllib.parse.urlparse(fn).path)
-      if ext not in ('', '.bz2'):
-        # old rlogs weren't bz2 compressed
+      if ext not in ('', '.bz2', '.zst'):
+        # old rlogs weren't compressed
         raise Exception(f"unknown extension {ext}")
 
       with FileReader(fn) as f:
@@ -43,6 +44,9 @@ class _LogFileReader:
 
     if ext == ".bz2" or dat.startswith(b'BZh9'):
       dat = bz2.decompress(dat)
+    elif ext == ".zst" or dat.startswith(b'\x28\xB5\x2F\xFD'):
+      # https://github.com/facebook/zstd/blob/dev/doc/zstd_compression_format.md#zstandard-frames
+      dat = zstd.decompress(dat)
 
     ents = capnp_log.Event.read_multiple_bytes(dat)