fix issues when transaction is unsampled in Kafka instrumentation

closes #1578

Author: beniwohli

elasticapm/instrumentation/packages/kafka.py

@@ -36,7 +36,7 @@
 from elasticapm.conf import constants
 from elasticapm.instrumentation.packages.base import AbstractInstrumentedModule
 from elasticapm.traces import DroppedSpan, capture_span, execution_context
-from elasticapm.utils.disttracing import TraceParent
+from elasticapm.utils.disttracing import TraceParent, TracingOptions
 
 
 class KafkaInstrumentation(AbstractInstrumentedModule):
@@ -48,6 +48,7 @@ class KafkaInstrumentation(AbstractInstrumentedModule):
  ]
  provider_name = "kafka"
  name = "kafka"
+ creates_transactions = True
 
  def _trace_send(self, instance, wrapped, *args, destination_info=None, **kwargs):
  topic = args[0] if args else kwargs["topic"]
@@ -68,7 +69,10 @@ def _trace_send(self, instance, wrapped, *args, destination_info=None, **kwargs)
  ) as span:
  transaction = execution_context.get_transaction()
  if transaction:
- tp = transaction.trace_parent.copy_from(span_id=span.id)
+ tp = transaction.trace_parent.copy_from(
+ span_id=span.id if span else transaction.id,
+ trace_options=None if span else TracingOptions(recorded=False),
+ )
  if headers:
  headers.append((constants.TRACEPARENT_BINARY_HEADER_NAME, tp.to_binary()))
  else:
@@ -79,22 +83,17 @@ def _trace_send(self, instance, wrapped, *args, destination_info=None, **kwargs)
  else:
  kwargs["headers"] = headers
  result = wrapped(*args, **kwargs)
- if instance and instance._metadata.controller and not isinstance(span, DroppedSpan):
+ if span and instance and instance._metadata.controller and not isinstance(span, DroppedSpan):
  address = instance._metadata.controller[1]
  port = instance._metadata.controller[2]
  span.context["destination"]["address"] = address
  span.context["destination"]["port"] = port
  return result
 
- def call_if_sampling(self, module, method, wrapped, instance, args, kwargs):
- # Contrasting to the superclass implementation, we *always* want to
- # return a proxied connection, even if there is no ongoing elasticapm
- # transaction yet. This ensures that we instrument the cursor once
- # the transaction started.
- return self.call(module, method, wrapped, instance, args, kwargs)
-
  def call(self, module, method, wrapped, instance, args, kwargs):
  client = get_client()
+ if client is None:
+ return wrapped(*args, **kwargs)
  destination_info = {
  "service": {"name": "kafka", "resource": "kafka/", "type": "messaging"},
  }
@@ -118,7 +117,7 @@ def call(self, module, method, wrapped, instance, args, kwargs):
  "destination": destination_info,
  },
  ) as span:
- if not isinstance(span, DroppedSpan) and instance._subscription.subscription:
+ if span and not isinstance(span, DroppedSpan) and instance._subscription.subscription:
  span.name += " from " + ", ".join(sorted(instance._subscription.subscription))
  results = wrapped(*args, **kwargs)
  return results
@@ -146,7 +145,7 @@ def call(self, module, method, wrapped, instance, args, kwargs):
  except StopIteration:
  span.cancel()
  raise
- if not isinstance(span, DroppedSpan):
+ if span and not isinstance(span, DroppedSpan):
  topic = result[0]
  if client.should_ignore_topic(topic):
  span.cancel()

tests/instrumentation/kafka_tests.py

@@ -207,3 +207,29 @@ def delayed_send():
  spans = elasticapm_client.events[SPAN]
  assert len(spans) == 1
  assert spans[0]["name"] == "Kafka POLL from bar, foo, test"
+
+
+def test_kafka_no_client(instrument, producer, consumer, topics):
+ assert elasticapm.get_client() is None
+ # the following code shouldn't trigger any errors
+ producer.send("test", key=b"foo", value=b"bar")
+ for item in consumer:
+ pass
+
+
+def test_kafka_send_unsampled_transaction(instrument, elasticapm_client, producer, topics):
+ transaction_object = elasticapm_client.begin_transaction("transaction")
+ transaction_object.is_sampled = False
+ producer.send("test", key=b"foo", value=b"bar")
+ elasticapm_client.end_transaction("foo")
+ spans = elasticapm_client.events[SPAN]
+ assert len(spans) == 0
+
+
+def test_kafka_poll_unsampled_transaction(instrument, elasticapm_client, consumer, topics):
+ transaction_object = elasticapm_client.begin_transaction("transaction")
+ transaction_object.is_sampled = False
+ consumer.poll(timeout_ms=50)
+ elasticapm_client.end_transaction("foo")
+ spans = elasticapm_client.events[SPAN]
+ assert len(spans) == 0