From 233ac8ef2a88c98fa08ca101bc96b9c0ef3f835c Mon Sep 17 00:00:00 2001 From: Mirochill Date: Mon, 25 May 2026 20:44:16 +0200 Subject: [PATCH] Decode raw V2 deflate histograms with headers --- hdrh/codec.py | 80 ++++++++++++++++++++++++++++++--------- test/test_hdrhistogram.py | 27 +++++++++++++ 2 files changed, 89 insertions(+), 18 deletions(-) diff --git a/hdrh/codec.py b/hdrh/codec.py index f9f380b..8695289 100644 --- a/hdrh/codec.py +++ b/hdrh/codec.py @@ -113,7 +113,7 @@ class HdrPayload(): '''A class that wraps the ctypes big endian struct that will hold the histogram wire format content (including the counters). ''' - def __init__(self, word_size, counts_len=0, compressed_payload=None): + def __init__(self, word_size, counts_len=0, compressed_payload=None, raw_payload=None): '''Two ways to use this class: - for an empty histogram (pass counts_len>0 and compressed_payload=None) - for a decoded histogram (counts_len=0 and compressed_payload!=None) @@ -128,6 +128,8 @@ def __init__(self, word_size, counts_len=0, compressed_payload=None): Caller must then invoke init_counts to pass in counts_len so that the counts array can be updated from the decoded varint buffer None if no compressed payload is to be associated to this instance + raw_payload (string) an uncompressed V2 payload that includes the + payload header and encoded counts ''' self.word_size = word_size self.counts_len = counts_len @@ -141,6 +143,8 @@ def __init__(self, word_size, counts_len=0, compressed_payload=None): raise ValueError('Invalid word size') if compressed_payload: self._decompress(compressed_payload) + elif raw_payload: + self._decode_payload(raw_payload) elif counts_len: self.payload = PayloadHeader() self.payload.cookie = get_encoding_cookie() @@ -169,16 +173,16 @@ def init_counts(self, counts_len): def get_counts(self): return self.counts - def _decompress(self, compressed_payload): - '''Decompress a compressed payload into this payload wrapper. - Note that the decompressed buffer is saved in self._data and the - counts array is not yet allocated. + def _decode_payload(self, payload_data): + '''Decode an uncompressed payload into this payload wrapper. + Note that the buffer is saved in self._data and the counts array is not + yet allocated. Args: - compressed_payload (string) a payload in zlib compressed form + payload_data (string) an uncompressed payload Exception: HdrCookieException: - the compressed payload has an invalid cookie + the payload has an invalid cookie HdrLengthException: the decompressed size is too small for the HdrPayload structure or is not aligned or is too large for the passed payload class @@ -188,10 +192,10 @@ def _decompress(self, compressed_payload): ''' # make sure this instance is pristine if self._data: - raise RuntimeError('Cannot decompress to an instance with payload') + raise RuntimeError('Cannot decode to an instance with payload') # Here it is important to keep a reference to the decompressed # string so that it does not get garbage collected - self._data = zlib.decompress(compressed_payload) + self._data = payload_data len_data = len(self._data) counts_size = len_data - payload_header_size @@ -208,6 +212,25 @@ def _decompress(self, compressed_payload): if word_size != V2_MAX_WORD_SIZE_IN_BYTES: raise HdrCookieException('Invalid V2 cookie: %x' % cookie) + def _decompress(self, compressed_payload): + '''Decompress a compressed payload into this payload wrapper. + Note that the decompressed buffer is saved in self._data and the + counts array is not yet allocated. + + Args: + compressed_payload (string) a payload in zlib compressed form + Exception: + HdrCookieException: + the compressed payload has an invalid cookie + HdrLengthException: + the decompressed size is too small for the HdrPayload structure + or is not aligned or is too large for the passed payload class + HdrHistogramSettingsException: + mismatch in the significant figures, lowest and highest + trackable value + ''' + self._decode_payload(zlib.decompress(compressed_payload)) + def compress(self, counts_limit): '''Compress this payload instance Args: @@ -339,18 +362,39 @@ def decode(encoded_histogram, b64_wrap=True): raise HdrLengthException('Base64 decoded message too short') header = ExternalHeader.from_buffer_copy(b64decode) - if get_cookie_base(header.cookie) != V2_COMPRESSION_COOKIE_BASE: + cookie_base = get_cookie_base(header.cookie) + if cookie_base == V2_COMPRESSION_COOKIE_BASE: + if header.length != b64dec_len - ext_header_size: + raise HdrLengthException('Decoded length=%d buffer length=%d' % + (header.length, b64dec_len - ext_header_size)) + # this will result in a copy of the compressed payload part + # could not find a way to do otherwise since zlib.decompress() + # expects a string (and does not like a buffer or a memoryview object) + cpayload = b64decode[ext_header_size:] + raw_payload = None + elif cookie_base == V2_ENCODING_COOKIE_BASE: + raw_payload = b64decode + cpayload = None + else: raise HdrCookieException() - if header.length != b64dec_len - ext_header_size: - raise HdrLengthException('Decoded length=%d buffer length=%d' % - (header.length, b64dec_len - ext_header_size)) - # this will result in a copy of the compressed payload part - # could not find a way to do otherwise since zlib.decompress() - # expects a string (and does not like a buffer or a memoryview object) - cpayload = b64decode[ext_header_size:] else: cpayload = encoded_histogram - hdr_payload = HdrPayload(8, compressed_payload=cpayload) + raw_payload = None + encoded_len = len(encoded_histogram) + if isinstance(encoded_histogram, (bytes, bytearray)) and encoded_len >= ext_header_size: + header = ExternalHeader.from_buffer_copy(encoded_histogram) + cookie_base = get_cookie_base(header.cookie) + if cookie_base == V2_COMPRESSION_COOKIE_BASE: + if header.length != encoded_len - ext_header_size: + raise HdrLengthException('Decoded length=%d buffer length=%d' % + (header.length, encoded_len - ext_header_size)) + cpayload = encoded_histogram[ext_header_size:] + elif cookie_base == V2_ENCODING_COOKIE_BASE: + raw_payload = encoded_histogram + if raw_payload: + hdr_payload = HdrPayload(8, raw_payload=raw_payload) + else: + hdr_payload = HdrPayload(8, compressed_payload=cpayload) return hdr_payload def add(self, other_encoder): diff --git a/test/test_hdrhistogram.py b/test/test_hdrhistogram.py index 01ceb2f..8ed0a0f 100644 --- a/test/test_hdrhistogram.py +++ b/test/test_hdrhistogram.py @@ -481,6 +481,33 @@ def test_hist_codec(): check_hist_codec_b64(word_size, True) check_hist_codec_b64(word_size, False) +@pytest.mark.codec +def test_hist_codec_raw_v2_deflate_header(): + rust_compressed = ( + b"\x1c\x84\x93\x14\x00\x00\x00\x1fx\x9c\x93i\x99," + b"\xcc\xc0\xc0\xc0\xcc\x00\x010\x9a\x11J3\xd9\x7f" + b"\x800\xfe32\x01\x00E\x0c\x03\x81" + ) + histogram = HdrHistogram.decode(rust_compressed, b64_wrap=False) + assert histogram.get_total_count() == 1 + +@pytest.mark.codec +def test_hist_codec_raw_v2_uncompressed(): + rust_uncompressed = ( + b"\x1c\x84\x93\x13\x00\x00\x00\x03\x00\x00\x00\x00" + b"\x00\x00\x00\x03\x00\x00\x00\x00\x00\x00\x00\x01" + b"\x00\x00\x00\x00\x00\x00\x00\x02?\xf0\x00\x00" + b"\x00\x00\x00\x00\xff\x01\x02" + ) + histogram = HdrHistogram.decode(rust_uncompressed, b64_wrap=False) + assert histogram.get_total_count() == 1 + +@pytest.mark.codec +def test_hist_codec_base64_v2_uncompressed(): + rust_uncompressed = "HISTEwAAAAMAAAAAAAAAAwAAAAAAAAABAAAAAAAAAAI/8AAAAAAAAP8BAg==" + histogram = HdrHistogram.decode(rust_uncompressed) + assert histogram.get_total_count() == 1 + @pytest.mark.codec def test_hist_codec_partial(): histogram = HdrHistogram(LOWEST, WRK2_MAX_LATENCY, SIGNIFICANT)