Decoding binary protocols in python

Decoding binary protocols like the MySQL Client/Server Protocol or MySQL's new X Protocol involves taking a sequence of bytes and turning them into integers.

In python the usually workhorse for this task is struct.unpack()

It takes a sequence of bytes and a format-string and returns a tuple of decoded values.

In the case of the MySQL Client/Server protocol the integers are (mostly) little-endian, unsigned and we can use:

format description
<B integer, little endian, unsigned, 1 bytes
<H integer, little endian, unsigned, 2 bytes
<L integer, little endian, unsigned, 4 bytes
<Q integer, little endian, unsigned, 8 bytes
# unpack_int_le_1

import struct


def unpack_int_le_struct(payload, l):
    if l == 0:
        return 0
    elif l == 1:
        return struct.unpack_from("<B", payload)[0]
    elif l == 2:
        return struct.unpack_from("<H", payload)[0]
    elif l == 4:
        return struct.unpack_from("<L", payload)[0]
    elif l == 8:
        return struct.unpack_from("<Q", payload)[0]
    else:
        # no native mapping for that byte-length,
        # fallback to shift+or
        v = 0
        sh = 0

        for d in struct.unpack_from("%dB" % l, payload):
            v |= d << sh
            sh += 8

        return v

The code gets benchmarked with 'timeit' which runs the setup once and the test-function 10000 times:

# benchmark.py

import timeit
import tabulate

from collections import OrderedDict

setup = r"""
from unpack_int_le_1 import unpack_int_le

test_str = b'\x00\x00\x00\x00\x00\x00\x00\x00'
"""


timings = OrderedDict()

for n in (1, 2, 3, 4, 8):
    key = str(n)

    timings[key] = timeit.timeit(
        r"unpack_int_le(test_str, %d)" % (n, ),
        setup=setup)


print(tabulate.tabulate(
    timings.items(),
    tablefmt="rst",
    headers=("byte length",
             "(s)",
             )))

shows us:

byte length time (s)
1 0.466556
2 0.422064
3 1.17804
4 0.439113
8 0.448069

As the MySQL Client/Server protocol also has a 3-byte unsigned integer which doesn't have mapping struct's format strings, the fallback is taken which is quite slow.

Optimizing 3 byte integers

There is a nice trick to speed it up:

uint3_le = struct.unpack("<L", payload[:3] + b"\x00")[0]
# unpack_int_le_2.py

import struct


def unpack_int_le(payload, l):
    if l == 0:
        return 0
    elif l == 1:
        return struct.unpack_from("<B", payload)[0]
    elif l == 2:
        return struct.unpack_from("<H", payload)[0]
    elif l == 3 and hasattr(payload, "__add__"):
        return struct.unpack_from("<L", payload[:l] + b"\x00")[0]
    elif l == 4:
        return struct.unpack_from("<L", payload)[0]
    elif l == 8:
        return struct.unpack_from("<Q", payload)[0]
    else:
        # no native mapping for that byte-length, fallback to
        # shift+or
        v = 0
        sh = 0

        for d in struct.unpack_from("%dB" % l, payload):
            v |= d << sh
            sh += 8

        return v

Note

the check for __add__ handles that case of memoryviews which can't be extended like that.

byte length time (s)
1 0.479504
2 0.430233
3 0.650036
4 0.456068
8 0.481057

The tiny bit of extra work to slice and extend the string, wins over doing all the work in pure python.

Naive cython

cython is a python to C converter which (with the help of a C compiler) builds native python modules.

You can install it via pip:

$ pip install cython

and generate a native python module with:

$ cythonize ./unpack_int_le_2.py && \
  gcc -shared -pthread -fPIC -fwrapv -O3 -Wall \
    -fno-strict-aliasing -I/usr/include/python2.7 \
    -o unpack_int_le_3.so unpack_int_le_3.c

Without any changes to the python code we gain between 20-40%:

byte length time (s)
1 0.36877
2 0.314847
3 0.427502
4 0.322987
8 0.334282

Static typed python with annotations

cython can do better though with static typing.

Looking at the generated unpack_int_le_3.c one can see that most of the time is spent in the abstractions (object, buffer-interface, ...) for each byte extracted.

cython can remove all that if we give it some hints:

# unpack_int_le_4.py

import cython


@cython.locals(payload=cython.p_uchar,
               l=cython.ulong,
               v=cython.ulong,
               sh=cython.ulong,
               n=cython.ulong)
def _unpack_int_le(payload, l):
    v = sh = 0

    for n in range(l):
        v |= payload[n] << sh
        sh += 8

    return v


@cython.locals(l=cython.ulong,
               charp=cython.p_uchar)
def unpack_int_le(payload, l):
    if l > len(payload):
        raise IndexError()

    if isinstance(payload, (bytes, bytearray)):
        charp = payload

        return _unpack_int_le(charp, l)
    else:
        raise TypeError(type(payload))

See also: http://docs.cython.org/en/latest/src/tutorial/pure.html

With these hints given the _unpack_int_le() from above gets translated to:

__pyx_v_v = 0;
__pyx_v_sh = 0;

__pyx_t_1 = __pyx_v_l;
for (__pyx_t_2 = 0; __pyx_t_2 < __pyx_t_1; __pyx_t_2+=1) {
  __pyx_v_n = __pyx_t_2;

  __pyx_v_v = (__pyx_v_v | ((__pyx_v_payload[__pyx_v_n]) << __pyx_v_sh));
  __pyx_v_sh = (__pyx_v_sh + 8);
}

return PyInt_from_unsigned_long(__pyx_v_v)

This gains another 50%:

byte length time (s)
1 0.224607
2 0.157248
3 0.158877
4 0.163004
8 0.165769

While the code is pure python with some annotations it sadly comes with some overhead in the generated C-code.

Static typed cython

cython has two more ways of specifying the static types:

  • move the static types into a .pxd file
  • write a .pyx file directly

In our case the result would be the same which is why I go with the first approach:

# unpack_int_le_5.pxd

import cython


@cython.locals(v=cython.ulong,
               sh=cython.ulong,
               n=cython.ulong)
cdef unsigned long _unpack_int_le(unsigned char *payload, unsigned long l)

@cython.locals(l=cython.ulong,
               charp=cython.p_uchar)
cpdef unpack_int_le(payload, unsigned long l)
# unpack_int_le_5.py

def _unpack_int_le(payload, l):
    v = sh = n = 0

    for n in range(l):
        v |= payload[n] << sh
        sh += 8

    return v


def unpack_int_le(payload, l):
    if l > len(payload):
        raise IndexError()

    if isinstance(payload, (bytes, bytearray)):
        charp = payload

        return _unpack_int_le(charp, l)
    else:
        raise TypeError(type(payload))
$ cythonize ./unpack_int_le_5.py && \
  gcc -shared -pthread -fPIC -fwrapv \
    -O3 -Wall -fno-strict-aliasing \
    -I/usr/include/python2.7 \
    -o unpack_int_le_5.so unpack_int_le_5.c
byte length time (s)
1 0.084414
2 0.0818369
3 0.083369
4 0.084229
8 0.0980709

One can see that the longer byte-length leads to slightly higher runtime.

Just cast it

One last step and we are at the final form of the optimizations which mirror what one would have written if one would have written the code in C directly.

If we are running on a little endian platform, we can save shift-or-loop and just cast the byte-sequence into the right integer directly.

Note

to show how the code looks in a pyx file, the previous py and pxd files have been merged into a pyx.

# unpack_int_le_6.pyx

import sys


cdef unsigned int is_le
is_le = sys.byteorder == "little"


cdef unsigned long _unpack_int_le(const unsigned char *payload, unsigned long l):
  cdef unsigned long *u32p
  cdef unsigned short *u16p
  cdef unsigned long long *u64p
  cdef unsigned long v = 0
  cdef unsigned long sh = 0
  cdef unsigned long n = 0

  if is_le:
      if l == 1:
          return payload[0]
      elif l == 2:
          u16p = <unsigned short *>payload
          return u16p[0]
      elif l == 4:
          u32p = <unsigned long *>payload
          return u32p[0]
      elif l == 8:
          u64p = <unsigned long long *>payload
          return u64p[0]

  v = sh = n = 0

  for n in range(l):
      v |= payload[n] << sh
      sh += 8

  return v


cpdef unsigned long unpack_int_le(payload, unsigned long l):
    cdef unsigned char *charp
    if l > len(payload):
        raise IndexError()

    if isinstance(payload, (bytes, bytearray)):
        charp = payload

        return _unpack_int_le(charp, l)
    else:
        raise TypeError(type(payload))
$ cythonize ./unpack_int_le_6.pyx && \
  gcc -shared -pthread -fPIC -fwrapv \
    -O3 -Wall -fno-strict-aliasing \
    -I/usr/include/python2.7 \
    -o unpack_int_le_6.so unpack_int_le_6.c
byte length time (s)
1 0.0812111
2 0.0792191
3 0.082288
4 0.079355
8 0.0791218

Note

the 3-byte case hits the loop again and is slightly slower.

Conclusion

Between the first attempt with struct.unpack() and our last with cython we have quite a speed difference:

byte length time (s) time (s) speedup
1 0.466556 0.0812111 5.75x
2 0.422064 0.0792191 5.34x
3 1.17804 0.082288 14.36x
4 0.439113 0.079355 5.56x
8 0.448069 0.0791218 5.67x

We got a 5x speed up and didn't had to write a single line of python C-API code.


Comments

Enable javascript to load comments.