Repository: dpkp/kafka-python
Pull request #2021 (2020-03-15 23:19:48)

Add support for `zstd` compression
The ZSTD relevant code, in a separate PR, @jeffwidman

<!-- Reviewable:start -->
---
This change is [<img src="https://reviewable.io/review_button.svg" height="34" align="absmiddle" alt="Reviewable"/>](https://reviewable.io/reviews/dpkp/kafka-python/2021)
<!-- Reviewable:end -->


jeffwidman: Also, in https://kafka-python.readthedocs.io/en/master/#compression can you mention support for `zstd`?

A clean way to do it would be to make it a bullet pointed list...
> kafka-python supports the following compression formats:
> * gzip  - support built-in 
> * lz4 - requires installing `python-lz4`
> * snappy - requires installing `python-snappy` (which requires the `snappy` `C` library)
> * zstd - requires installing `zstandard`


jeffwidman: Now that https://github.com/dpkp/kafka-python/pull/2020 is merged, can you please rebase this on top of that? 
gabriel-tincu: > Now that #2020 is merged, can you please rebase this on top of that?

done, though it seems github diff algorithm or needs to refresh master the branch tip
jeffwidman: The diff is appearing fine for me.

Typically what you'd do in this case is `git rebase -i origin/master` and then squash all these commits into a single one...or just a couple of logically separated ones... ie, the fixup commits get merged together.
jeffwidman: Actually strike that, in this diff I'm suddenly seeing the `log_start_offset` stuff that has already been merged...

Are you sure you pulled latest `master` and then rebased on top of that?
gabriel-tincu: > Actually strike that, in this diff I'm suddenly seeing the `log_start_offset` stuff that has already been merged...
> 
> Are you sure you pulled latest `master` and then rebased on top of that?

i rebased on top of the merged branch, due to both PR's being made from a forked version of the repo. That is to say, that repo's master is not in sync with this repo's master, for various reasons. 
That is why i linked the compare link, which compares the tips of both branches (or rather tip-1 for the source branch, as i made a commit just a few minutes ago), to illustrate that using compare, the diff appears as it should
jeffwidman: > i rebased on top of the merged branch, due to both PR's being made from a forked version of the repo. That is to say, that repo's master is not in sync with this repo's master, for various reasons.
That is why i linked the compare link, which compares the tips of both branches (or rather tip-1 for the source branch, as i made a commit just a few minutes ago), to illustrate that using compare, the diff appears as it should

Sorry, I wasn't clear... you don't want to rebase on top of the merged PR branch, you want to rebase on top of `master` of this repo, which contains the merged branch...

We'll need to see the diff so we can see _exactly_ what is getting merged. 
gabriel-tincu: > > i rebased on top of the merged branch, due to both PR's being made from a forked version of the repo. That is to say, that repo's master is not in sync with this repo's master, for various reasons.
> > That is why i linked the compare link, which compares the tips of both branches (or rather tip-1 for the source branch, as i made a commit just a few minutes ago), to illustrate that using compare, the diff appears as it should
> 
> Sorry, I wasn't clear... you don't want to rebase on top of the merged PR branch, you want to rebase on top of `master` of this repo, which contains the merged branch...
> 
> We'll need to see the diff so we can see _exactly_ what is getting merged.

For future reference, that would imply updating the origin first, then doing the actual rebase, or just adding a new remote, correct? I'll try to open a new PR with an identical but differently named branch 
jeffwidman: Can you just merge `master` into this one? It should fix the diff w/o needing to rebase or anything...
gabriel-tincu: > Can you just pull latest `master` from this repo and then merge that `master` into this branch? It should fix the diff w/o needing to rebase or anything...
> 
> And then I'll cleanup the history when I squash merge. If this truly doesn't work, then it's a weird github bug, but I expect it should work since we often do that workflow at my day job with no issues.

That seems to have worked
jeffwidman: IMO this is mostly ready to go, except for https://github.com/dpkp/kafka-python/pull/2021#discussion_r398247782... fixing the root issue of that is outside the scope of this particular PR, but in the meantime I'd rather not add that workaround hack.

It looks like #2024 takes a stab at fixing the root cause, let's see if we can land that and then land this...
Green-Angry-Bird: Our team is looking for exactly this. When can we merge to master. :D 
tvoinarovskyi: @gabriel-tincu Sorry, but could you rebase this over master now? Thanks =) 
gabriel-tincu: > @gabriel-tincu Sorry, but could you rebase this over master now? Thanks =)

@tvoinarovskyi Done. I noticed that tox fails specifically for python3.7, but with an error related to tox itself rather than anything related to this PR
Green-Angry-Bird: 👍 I can't wait for my consumer to not crash on zstd encoded messages.
Green-Angry-Bird: Is there any update on this?  Our team is looking forward to using this library with our zstd encoded topics.

Patch:
diff --git a/.travis.yml b/.travis.yml
index b98aa16..e837924 100644
--- a/.travis.yml
+++ b/.travis.yml
@@ -6,6 +6,7 @@ python:
     - 2.7
     - 3.4
     - 3.7
+    - 3.8
     - pypy2.7-6.0
 
 env:
@@ -21,6 +22,7 @@ addons:
   apt:
     packages:
       - libsnappy-dev
+      - libzstd-dev
       - openjdk-8-jdk
 
 cache:
diff --git a/README.rst b/README.rst
index bae567b..29e6935 100644
--- a/README.rst
+++ b/README.rst
@@ -2,7 +2,7 @@ Kafka Python client
 ------------------------
 
 .. image:: https://img.shields.io/badge/kafka-2.4%2C%202.3%2C%202.2%2C%202.1%2C%202.0%2C%201.1%2C%201.0%2C%200.11%2C%200.10%2C%200.9%2C%200.8-brightgreen.svg
-    :target: https://kafka-python.readthedocs.io/compatibility.html
+    :target: https://kafka-python.readthedocs.io/en/master/compatibility.html
 .. image:: https://img.shields.io/pypi/pyversions/kafka-python.svg
     :target: https://pypi.python.org/pypi/kafka-python
 .. image:: https://coveralls.io/repos/dpkp/kafka-python/badge.svg?branch=master&service=github
diff --git a/benchmarks/README b/benchmarks/README.md
similarity index 100%
rename from benchmarks/README
rename to benchmarks/README.md
diff --git a/docs/compatibility.rst b/docs/compatibility.rst
index 93be6fd..ae15261 100644
--- a/docs/compatibility.rst
+++ b/docs/compatibility.rst
@@ -16,6 +16,6 @@ Although kafka-python is tested and expected to work on recent broker versions,
 not all features are supported. Specifically, authentication codecs, and
 transactional producer/consumer support are not fully implemented. PRs welcome!
 
-kafka-python is tested on python 2.7, 3.4, 3.7, and pypy2.7.
+kafka-python is tested on python 2.7, 3.4, 3.7, 3.8 and pypy2.7.
 
 Builds and tests via Travis-CI.  See https://travis-ci.org/dpkp/kafka-python
diff --git a/docs/conf.py b/docs/conf.py
index c7da0bc..efa8d08 100644
--- a/docs/conf.py
+++ b/docs/conf.py
@@ -12,7 +12,6 @@
 # All configuration values have a default; values that are commented out
 # serve to show the default.
 
-import sys
 import os
 
 # If extensions (or modules to document with autodoc) are in another directory,
diff --git a/docs/index.rst b/docs/index.rst
index fa6f93c..9c46e33 100644
--- a/docs/index.rst
+++ b/docs/index.rst
@@ -122,11 +122,12 @@ multiprocessing is recommended.
 Compression
 ***********
 
-kafka-python supports gzip compression/decompression natively. To produce or
-consume lz4 compressed messages, you should install python-lz4 (pip install lz4).
-To enable snappy, install python-snappy (also requires snappy library).
-See `Installation <install.html#optional-snappy-install>`_ for more information.
+kafka-python supports multiple compression types:
 
+ - gzip : supported natively
+ - lz4 : requires `python-lz4 <https://pypi.org/project/lz4/>`_ installed
+ - snappy : requires the `python-snappy <https://pypi.org/project/python-snappy/>`_  package (which requires the snappy C library)
+ - zstd : requires the `python-zstandard <https://github.com/indygreg/python-zstandard>`_ package installed
 
 Protocol
 ********
diff --git a/kafka/client_async.py b/kafka/client_async.py
index 5379153..caa88cf 100644
--- a/kafka/client_async.py
+++ b/kafka/client_async.py
@@ -2,7 +2,6 @@ from __future__ import absolute_import, division
 
 import collections
 import copy
-import functools
 import logging
 import random
 import socket
diff --git a/kafka/codec.py b/kafka/codec.py
index aa9fc82..917400e 100644
--- a/kafka/codec.py
+++ b/kafka/codec.py
@@ -10,12 +10,18 @@ from kafka.vendor.six.moves import range
 
 _XERIAL_V1_HEADER = (-126, b'S', b'N', b'A', b'P', b'P', b'Y', 0, 1, 1)
 _XERIAL_V1_FORMAT = 'bccccccBii'
+ZSTD_MAX_OUTPUT_SIZE = 1024 * 1024
 
 try:
     import snappy
 except ImportError:
     snappy = None
 
+try:
+    import zstandard as zstd
+except ImportError:
+    zstd = None
+
 try:
     import lz4.frame as lz4
 
@@ -58,6 +64,10 @@ def has_snappy():
     return snappy is not None
 
 
+def has_zstd():
+    return zstd is not None
+
+
 def has_lz4():
     if lz4 is not None:
         return True
@@ -299,3 +309,18 @@ def lz4_decode_old_kafka(payload):
         payload[header_size:]
     ])
     return lz4_decode(munged_payload)
+
+
+def zstd_encode(payload):
+    if not zstd:
+        raise NotImplementedError("Zstd codec is not available")
+    return zstd.ZstdCompressor().compress(payload)
+
+
+def zstd_decode(payload):
+    if not zstd:
+        raise NotImplementedError("Zstd codec is not available")
+    try:
+        return zstd.ZstdDecompressor().decompress(payload)
+    except zstd.ZstdError:
+        return zstd.ZstdDecompressor().decompress(payload, max_output_size=ZSTD_MAX_OUTPUT_SIZE)
diff --git a/kafka/consumer/fetcher.py b/kafka/consumer/fetcher.py
index 2c11eb9..e4f8c18 100644
--- a/kafka/consumer/fetcher.py
+++ b/kafka/consumer/fetcher.py
@@ -293,7 +293,7 @@ class Fetcher(six.Iterator):
                 # Issue #1780
                 # Recheck partition existence after after a successful metadata refresh
                 if refresh_future.succeeded() and isinstance(future.exception, Errors.StaleMetadata):
-                    log.debug("Stale metadata was raised, and we now have an updated metadata. Rechecking partition existance")
+                    log.debug("Stale metadata was raised, and we now have an updated metadata. Rechecking partition existence")
                     unknown_partition = future.exception.args[0]  # TopicPartition from StaleMetadata
                     if self._client.cluster.leader_for_partition(unknown_partition) is None:
                         log.debug("Removed partition %s from offsets retrieval" % (unknown_partition, ))
diff --git a/kafka/coordinator/base.py b/kafka/coordinator/base.py
index b0e236a..5e41309 100644
--- a/kafka/coordinator/base.py
+++ b/kafka/coordinator/base.py
@@ -3,7 +3,6 @@ from __future__ import absolute_import, division
 import abc
 import copy
 import logging
-import sys
 import threading
 import time
 import weakref
@@ -243,7 +242,7 @@ class BaseCoordinator(object):
         """Block until the coordinator for this group is known
         (and we have an active connection -- java client uses unsent queue).
         """
-        with self._lock:
+        with self._client._lock, self._lock:
             while self.coordinator_unknown():
 
                 # Prior to 0.8.2 there was no group coordinator
@@ -346,7 +345,7 @@ class BaseCoordinator(object):
 
     def ensure_active_group(self):
         """Ensure that the group is active (i.e. joined and synced)"""
-        with self._lock:
+        with self._client._lock, self._lock:
             if self._heartbeat_thread is None:
                 self._start_heartbeat_thread()
 
@@ -764,7 +763,7 @@ class BaseCoordinator(object):
 
     def maybe_leave_group(self):
         """Leave the current group and reset local generation/memberId."""
-        with self._lock:
+        with self._client._lock, self._lock:
             if (not self.coordinator_unknown()
                 and self.state is not MemberState.UNJOINED
                 and self._generation is not Generation.NO_GENERATION):
@@ -947,6 +946,15 @@ class HeartbeatThread(threading.Thread):
             log.debug('Heartbeat thread closed')
 
     def _run_once(self):
+        with self.coordinator._client._lock, self.coordinator._lock:
+            if self.enabled and self.coordinator.state is MemberState.STABLE:
+                # TODO: When consumer.wakeup() is implemented, we need to
+                # disable here to prevent propagating an exception to this
+                # heartbeat thread
+                # must get client._lock, or maybe deadlock at heartbeat 
+                # failure callbak in consumer poll
+                self.coordinator._client.poll(timeout_ms=0)
+
         with self.coordinator._lock:
             if not self.enabled:
                 log.debug('Heartbeat disabled. Waiting')
@@ -962,11 +970,6 @@ class HeartbeatThread(threading.Thread):
                 self.disable()
                 return
 
-            # TODO: When consumer.wakeup() is implemented, we need to
-            # disable here to prevent propagating an exception to this
-            # heartbeat thread
-            self.coordinator._client.poll(timeout_ms=0)
-
             if self.coordinator.coordinator_unknown():
                 future = self.coordinator.lookup_coordinator()
                 if not future.is_done or future.failed():
diff --git a/kafka/producer/kafka.py b/kafka/producer/kafka.py
index 9509ab9..dba1801 100644
--- a/kafka/producer/kafka.py
+++ b/kafka/producer/kafka.py
@@ -12,7 +12,7 @@ from kafka.vendor import six
 
 import kafka.errors as Errors
 from kafka.client_async import KafkaClient, selectors
-from kafka.codec import has_gzip, has_snappy, has_lz4
+from kafka.codec import has_gzip, has_snappy, has_lz4, has_zstd
 from kafka.metrics import MetricConfig, Metrics
 from kafka.partitioner.default import DefaultPartitioner
 from kafka.producer.future import FutureRecordMetadata, FutureProduceResult
@@ -119,7 +119,7 @@ class KafkaProducer(object):
                 available guarantee.
             If unset, defaults to acks=1.
         compression_type (str): The compression type for all data generated by
-            the producer. Valid values are 'gzip', 'snappy', 'lz4', or None.
+            the producer. Valid values are 'gzip', 'snappy', 'lz4', 'zstd' or None.
             Compression is of full batches of data, so the efficacy of batching
             will also impact the compression ratio (more batching means better
             compression). Default: None.
@@ -339,6 +339,7 @@ class KafkaProducer(object):
         'gzip': (has_gzip, LegacyRecordBatchBuilder.CODEC_GZIP),
         'snappy': (has_snappy, LegacyRecordBatchBuilder.CODEC_SNAPPY),
         'lz4': (has_lz4, LegacyRecordBatchBuilder.CODEC_LZ4),
+        'zstd': (has_zstd, DefaultRecordBatchBuilder.CODEC_ZSTD),
         None: (lambda: True, LegacyRecordBatchBuilder.CODEC_NONE),
     }
 
@@ -388,6 +389,9 @@ class KafkaProducer(object):
         if self.config['compression_type'] == 'lz4':
             assert self.config['api_version'] >= (0, 8, 2), 'LZ4 Requires >= Kafka 0.8.2 Brokers'
 
+        if self.config['compression_type'] == 'zstd':
+            assert self.config['api_version'] >= (2, 1, 0), 'Zstd Requires >= Kafka 2.1.0 Brokers'
+
         # Check compression_type for library support
         ct = self.config['compression_type']
         if ct not in self._COMPRESSORS:
diff --git a/kafka/protocol/message.py b/kafka/protocol/message.py
index 31527bf..4c5c031 100644
--- a/kafka/protocol/message.py
+++ b/kafka/protocol/message.py
@@ -3,8 +3,8 @@ from __future__ import absolute_import
 import io
 import time
 
-from kafka.codec import (has_gzip, has_snappy, has_lz4,
-                     gzip_decode, snappy_decode,
+from kafka.codec import (has_gzip, has_snappy, has_lz4, has_zstd,
+                     gzip_decode, snappy_decode, zstd_decode,
                      lz4_decode, lz4_decode_old_kafka)
 from kafka.protocol.frame import KafkaBytes
 from kafka.protocol.struct import Struct
@@ -35,6 +35,7 @@ class Message(Struct):
     CODEC_GZIP = 0x01
     CODEC_SNAPPY = 0x02
     CODEC_LZ4 = 0x03
+    CODEC_ZSTD = 0x04
     TIMESTAMP_TYPE_MASK = 0x08
     HEADER_SIZE = 22  # crc(4), magic(1), attributes(1), timestamp(8), key+value size(4*2)
 
@@ -119,7 +120,7 @@ class Message(Struct):
 
     def decompress(self):
         codec = self.attributes & self.CODEC_MASK
-        assert codec in (self.CODEC_GZIP, self.CODEC_SNAPPY, self.CODEC_LZ4)
+        assert codec in (self.CODEC_GZIP, self.CODEC_SNAPPY, self.CODEC_LZ4, self.CODEC_ZSTD)
         if codec == self.CODEC_GZIP:
             assert has_gzip(), 'Gzip decompression unsupported'
             raw_bytes = gzip_decode(self.value)
@@ -132,6 +133,9 @@ class Message(Struct):
                 raw_bytes = lz4_decode_old_kafka(self.value)
             else:
                 raw_bytes = lz4_decode(self.value)
+        elif codec == self.CODEC_ZSTD:
+            assert has_zstd(), "ZSTD decompression unsupported"
+            raw_bytes = zstd_decode(self.value)
         else:
             raise Exception('This should be impossible')
 
diff --git a/kafka/record/default_records.py b/kafka/record/default_records.py
index 07368bb..a098c42 100644
--- a/kafka/record/default_records.py
+++ b/kafka/record/default_records.py
@@ -62,8 +62,8 @@ from kafka.record.util import (
 )
 from kafka.errors import CorruptRecordException, UnsupportedCodecError
 from kafka.codec import (
-    gzip_encode, snappy_encode, lz4_encode,
-    gzip_decode, snappy_decode, lz4_decode
+    gzip_encode, snappy_encode, lz4_encode, zstd_encode,
+    gzip_decode, snappy_decode, lz4_decode, zstd_decode
 )
 import kafka.codec as codecs
 
@@ -97,6 +97,7 @@ class DefaultRecordBase(object):
     CODEC_GZIP = 0x01
     CODEC_SNAPPY = 0x02
     CODEC_LZ4 = 0x03
+    CODEC_ZSTD = 0x04
     TIMESTAMP_TYPE_MASK = 0x08
     TRANSACTIONAL_MASK = 0x10
     CONTROL_MASK = 0x20
@@ -111,6 +112,8 @@ class DefaultRecordBase(object):
             checker, name = codecs.has_snappy, "snappy"
         elif compression_type == self.CODEC_LZ4:
             checker, name = codecs.has_lz4, "lz4"
+        elif compression_type == self.CODEC_ZSTD:
+            checker, name = codecs.has_zstd, "zstd"
         if not checker():
             raise UnsupportedCodecError(
                 "Libraries for {} compression codec not found".format(name))
@@ -185,6 +188,8 @@ class DefaultRecordBatch(DefaultRecordBase, ABCRecordBatch):
                     uncompressed = snappy_decode(data.tobytes())
                 if compression_type == self.CODEC_LZ4:
                     uncompressed = lz4_decode(data.tobytes())
+                if compression_type == self.CODEC_ZSTD:
+                    uncompressed = zstd_decode(data.tobytes())
                 self._buffer = bytearray(uncompressed)
                 self._pos = 0
         self._decompressed = True
@@ -517,6 +522,8 @@ class DefaultRecordBatchBuilder(DefaultRecordBase, ABCRecordBatchBuilder):
                 compressed = snappy_encode(data)
             elif self._compression_type == self.CODEC_LZ4:
                 compressed = lz4_encode(data)
+            elif self._compression_type == self.CODEC_ZSTD:
+                compressed = zstd_encode(data)
             compressed_size = len(compressed)
             if len(data) <= compressed_size:
                 # We did not get any benefit from compression, lets send
diff --git a/kafka/record/memory_records.py b/kafka/record/memory_records.py
index a6c4b51..fc2ef2d 100644
--- a/kafka/record/memory_records.py
+++ b/kafka/record/memory_records.py
@@ -117,7 +117,7 @@ class MemoryRecordsBuilder(object):
 
     def __init__(self, magic, compression_type, batch_size):
         assert magic in [0, 1, 2], "Not supported magic"
-        assert compression_type in [0, 1, 2, 3], "Not valid compression type"
+        assert compression_type in [0, 1, 2, 3, 4], "Not valid compression type"
         if magic >= 2:
             self._builder = DefaultRecordBatchBuilder(
                 magic=magic, compression_type=compression_type,
diff --git a/kafka/scram.py b/kafka/scram.py
index 684925c..7f00375 100644
--- a/kafka/scram.py
+++ b/kafka/scram.py
@@ -71,7 +71,6 @@ class ScramClient:
         )
 
     def final_message(self):
-        client_final_no_proof = 'c=biws,r=' + self.nonce
         return 'c=biws,r={},p={}'.format(self.nonce, base64.b64encode(self.client_proof).decode('utf-8'))
 
     def process_server_final_message(self, server_final_message):
diff --git a/requirements-dev.txt b/requirements-dev.txt
index d283090..6c3a666 100644
--- a/requirements-dev.txt
+++ b/requirements-dev.txt
@@ -1,6 +1,6 @@
 flake8==3.4.1
-pytest==3.10.0
-pytest-cov==2.6.0
+pytest==5.4.3
+pytest-cov==2.10.0
 docker-py==1.10.6
 coveralls==1.5.1
 Sphinx==1.6.4
@@ -9,8 +9,8 @@ xxhash==1.3.0
 python-snappy==0.5.3
 tox==3.5.3
 mock==3.0.5
-pylint==1.9.3
-pytest-pylint==0.12.3
+pylint==2.5.3
+pytest-pylint==0.17.0
 pytest-mock==1.10.0
 sphinx-rtd-theme==0.2.4
 crc32c==1.7
diff --git a/setup.py b/setup.py
index 8bc484c..005c5ad 100644
--- a/setup.py
+++ b/setup.py
@@ -56,6 +56,7 @@ setup(
         "Programming Language :: Python :: 3.5",
         "Programming Language :: Python :: 3.6",
         "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
         "Programming Language :: Python :: Implementation :: PyPy",
         "Topic :: Software Development :: Libraries :: Python Modules",
     ]
diff --git a/tox.ini b/tox.ini
index 06403d6..10e9911 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,5 +1,5 @@
 [tox]
-envlist = py{26,27,34,35,36,37,py}, docs
+envlist = py{26,27,34,35,36,37,38,py}, docs
 
 [pytest]
 testpaths = kafka test
@@ -8,13 +8,14 @@ log_format = %(created)f %(filename)-23s %(threadName)s %(message)s
 
 [testenv]
 deps =
-    pytest<4.0
+    pytest
     pytest-cov
-    py{27,34,35,36,37,py}: pylint
-    py{27,34,35,36,37,py}: pytest-pylint
+    py{27,34,35,36,37,38,py}: pylint
+    py{27,34,35,36,37,38,py}: pytest-pylint
     pytest-mock
     mock
     python-snappy
+    zstandard
     lz4
     xxhash
     crc32c


Test patch:
diff --git a/test/test_codec.py b/test/test_codec.py
index 9eff888..e057074 100644
--- a/test/test_codec.py
+++ b/test/test_codec.py
@@ -7,11 +7,12 @@ import pytest
 from kafka.vendor.six.moves import range
 
 from kafka.codec import (
-    has_snappy, has_lz4,
+    has_snappy, has_lz4, has_zstd,
     gzip_encode, gzip_decode,
     snappy_encode, snappy_decode,
     lz4_encode, lz4_decode,
     lz4_encode_old_kafka, lz4_decode_old_kafka,
+    zstd_encode, zstd_decode,
 )
 
 from test.testutil import random_string
@@ -113,3 +114,11 @@ def test_lz4_incremental():
         b2 = lz4_decode(lz4_encode(b1))
         assert len(b1) == len(b2)
         assert b1 == b2
+
+
+@pytest.mark.skipif(not has_zstd(), reason="Zstd not available")
+def test_zstd():
+    for _ in range(1000):
+        b1 = random_string(100).encode('utf-8')
+        b2 = zstd_decode(zstd_encode(b1))
+        assert b1 == b2
diff --git a/test/test_producer.py b/test/test_producer.py
index 9605adf..7263130 100644
--- a/test/test_producer.py
+++ b/test/test_producer.py
@@ -23,16 +23,16 @@ def test_buffer_pool():
 
 
 @pytest.mark.skipif(not env_kafka_version(), reason="No KAFKA_VERSION set")
-@pytest.mark.parametrize("compression", [None, 'gzip', 'snappy', 'lz4'])
+@pytest.mark.parametrize("compression", [None, 'gzip', 'snappy', 'lz4', 'zstd'])
 def test_end_to_end(kafka_broker, compression):
-
     if compression == 'lz4':
-        # LZ4 requires 0.8.2
         if env_kafka_version() < (0, 8, 2):
-            return
-        # python-lz4 crashes on older versions of pypy
+            pytest.skip('LZ4 requires 0.8.2')
         elif platform.python_implementation() == 'PyPy':
-            return
+            pytest.skip('python-lz4 crashes on older versions of pypy')
+
+    if compression == 'zstd' and env_kafka_version() < (2, 1, 0):
+        pytest.skip('zstd requires kafka 2.1.0 or newer')
 
     connect_str = ':'.join([kafka_broker.host, str(kafka_broker.port)])
     producer = KafkaProducer(bootstrap_servers=connect_str,
@@ -81,8 +81,10 @@ def test_kafka_producer_gc_cleanup():
 
 
 @pytest.mark.skipif(not env_kafka_version(), reason="No KAFKA_VERSION set")
-@pytest.mark.parametrize("compression", [None, 'gzip', 'snappy', 'lz4'])
+@pytest.mark.parametrize("compression", [None, 'gzip', 'snappy', 'lz4', 'zstd'])
 def test_kafka_producer_proper_record_metadata(kafka_broker, compression):
+    if compression == 'zstd' and env_kafka_version() < (2, 1, 0):
+        pytest.skip('zstd requires 2.1.0 or more')
     connect_str = ':'.join([kafka_broker.host, str(kafka_broker.port)])
     producer = KafkaProducer(bootstrap_servers=connect_str,
                              retries=5,
@@ -124,10 +126,8 @@ def test_kafka_producer_proper_record_metadata(kafka_broker, compression):
     if headers:
         assert record.serialized_header_size == 22
 
-    # generated timestamp case is skipped for broker 0.9 and below
     if magic == 0:
-        return
-
+        pytest.skip('generated timestamp case is skipped for broker 0.9 and below')
     send_time = time.time() * 1000
     future = producer.send(
         topic,