Kafka Producer Instrumentation (#612)

* Basic kafka producer instrumentation.
Co-authored-by: Lalleh Rafeei <[email protected]>

* Add testing for kafka producer.
Co-authored-by: Lalleh Rafeei <[email protected]>

* Add producer test for notice_error & cleanup

Co-authored-by: Lalleh Rafeei <[email protected]>

* Fix py2 naming

* Fix lint error

* Kafka Heartbeat (#614)

* Add Kafka test infra.

* Update tests.yml.

* Basic kafka producer instrumentation.
Co-authored-by: Lalleh Rafeei <[email protected]>

* Add testing for kafka producer.
Co-authored-by: Lalleh Rafeei <[email protected]>

* Kafka Heartbeat instrumentation

* Remove changes not related to heartbeat

Co-authored-by: Uma Annamalai <[email protected]>
Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>

Co-authored-by: Hannah Stepanek <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>
Co-authored-by: Uma Annamalai <[email protected]>
Co-authored-by: Lalleh Rafeei <[email protected]>

Author: 4 people

newrelic/config.py

@@ -2307,6 +2307,11 @@ def _process_module_builtin_defaults():
         "instrument_cherrypy__cptree",
     )
 
+    _process_module_definition(
+        "kafka.producer.kafka",
+        "newrelic.hooks.messagebroker_kafkapython",
+        "instrument_kafka_producer",
+    )
     _process_module_definition(
         "kafka.coordinator.heartbeat",
         "newrelic.hooks.messagebroker_kafkapython",

newrelic/hooks/messagebroker_kafkapython.py

@@ -13,6 +13,9 @@
 # limitations under the License.
 
 from newrelic.api.application import application_instance
+from newrelic.api.message_trace import MessageTrace
+from newrelic.api.time_trace import notice_error
+from newrelic.api.transaction import current_transaction
 from newrelic.common.object_wrapper import wrap_function_wrapper
 
 HEARTBEAT_POLL = "MessageBroker/Kafka/Heartbeat/Poll"
@@ -23,6 +26,40 @@
 HEARTBEAT_POLL_TIMEOUT = "MessageBroker/Kafka/Heartbeat/PollTimeout"
 
 
+def _bind_send(topic, value=None, key=None, headers=None, partition=None, timestamp_ms=None):
+    return topic, value, key, headers, partition, timestamp_ms
+
+
+def wrap_KafkaProducer_send(wrapped, instance, args, kwargs):
+    transaction = current_transaction()
+
+    if transaction is None:
+        return wrapped(*args, **kwargs)
+
+    topic, value, key, headers, partition, timestamp_ms = _bind_send(*args, **kwargs)
+    headers = list(headers) if headers else []
+
+    with MessageTrace(
+        library="Kafka",
+        operation="Produce",
+        destination_type="Topic",
+        destination_name=topic or "Default",
+        source=wrapped,
+    ) as trace:
+        dt_headers = [(k, v.encode("utf-8")) for k, v in trace.generate_request_headers(transaction)]
+        headers.extend(dt_headers)
+        try:
+            return wrapped(topic, value=value, key=key, headers=headers, partition=partition, timestamp_ms=timestamp_ms)
+        except Exception:
+            notice_error()
+            raise
+
+
+def instrument_kafka_producer(module):
+    if hasattr(module, "KafkaProducer"):
+        wrap_function_wrapper(module, "KafkaProducer.send", wrap_KafkaProducer_send)
+
+
 def metric_wrapper(metric_name, check_result=False):
     def _metric_wrapper(wrapped, instance, args, kwargs):
         result = wrapped(*args, **kwargs)

tests/messagebroker_kafkapython/test_kafka_produce.py

@@ -12,19 +12,64 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import time
+import pytest
+from conftest import cache_kafka_headers
+from testing_support.fixtures import (
+    validate_non_transaction_error_event,
+    validate_transaction_metrics,
+)
+from testing_support.validators.validate_messagebroker_headers import (
+    validate_messagebroker_headers,
+)
 
+from newrelic.api.background_task import background_task
+from newrelic.packages import six
 
-def test_no_harm(topic, producer, consumer):
-    MESSAGES = [
-        {"foo": "bar"},
-        {"baz": "bat"},
-    ]
 
-    for msg in MESSAGES:
-        time.sleep(1)
-        producer.send(topic, value=msg)
-    producer.flush()
+def test_producer_records_trace(topic, send_producer_messages):
+    scoped_metrics = [("MessageBroker/Kafka/Topic/Produce/Named/%s" % topic, 3)]
+    unscoped_metrics = scoped_metrics
+    txn_name = "test_kafka_produce:test_producer_records_trace.<locals>.test" if six.PY3 else "test_kafka_produce:test"
 
-    for msg in consumer:
-        assert msg.topic == topic
+    @validate_transaction_metrics(
+        txn_name,
+        scoped_metrics=scoped_metrics,
+        rollup_metrics=unscoped_metrics,
+        background_task=True,
+    )
+    @background_task()
+    @cache_kafka_headers
+    @validate_messagebroker_headers
+    def test():
+        send_producer_messages()
+
+    test()
+
+
+def test_producer_records_error_if_raised(topic, producer):
+    _intrinsic_attributes = {
+        "error.class": "AssertionError",
+        "error.message": "Need at least one: key or value",
+        "error.expected": False,
+    }
+
+    @validate_non_transaction_error_event(_intrinsic_attributes)
+    @background_task()
+    def test():
+        producer.send(topic, None)
+        producer.flush()
+
+    with pytest.raises(AssertionError):
+        test()
+
+
+@pytest.fixture
+def send_producer_messages(topic, producer):
+    def _test():
+        messages = [1, 2, 3]
+        for message in messages:
+            producer.send(topic, message)
+
+        producer.flush()
+
+    return _test