Skip to content

Decode raw Rust V2 histograms #59

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
Mirochill wants to merge 1 commit into
base: master
Choose a base branch
from
Draft

Decode raw Rust V2 histograms #59

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
80 changes: 62 additions & 18 deletions hdrh/codec.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand All @@ -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
Expand All @@ -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()
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand All @@ -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:
Expand Down Expand Up @@ -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):
Expand Down
27 changes: 27 additions & 0 deletions test/test_hdrhistogram.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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)
Expand Down