29
29
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
30
30
31
31
import time
32
+ from typing import Optional
32
33
33
34
import elasticapm
35
+ from elasticapm import get_client
34
36
from elasticapm .conf import constants
35
- from elasticapm .contrib .django .client import get_client as d_client
36
- from elasticapm .contrib .flask import get_client as f_client
37
37
from elasticapm .instrumentation .packages .base import AbstractInstrumentedModule
38
38
from elasticapm .traces import capture_span , execution_context
39
39
from elasticapm .utils .disttracing import TraceParent
40
40
41
41
42
- def django_client ():
43
- _client = None
44
- try :
45
- _client = d_client ()
46
- except Exception : # in some cases we get a different exception
47
- return None
48
- return _client
49
-
50
-
51
- def get_client ():
52
- if django_client ():
53
- return django_client ()
54
- elif f_client ():
55
- return f_client ()
56
-
57
-
58
- def get_trace_id (result ):
59
- for i in result :
60
- if isinstance (i , list ):
61
- if isinstance (i [0 ], tuple ) and len (i [0 ]) == 2 :
62
- for k , v in i :
63
- k_str = str (k )
64
- if k_str == "trace" :
65
- return v
66
- return None
67
-
68
-
69
42
class KafkaInstrumentation (AbstractInstrumentedModule ):
70
43
71
44
instrument_list = [
72
45
("kafka" , "KafkaProducer.send" ),
73
- ("kafka" , "KafkaConsumer.next" ),
46
+ ("kafka" , "KafkaConsumer.poll" ),
47
+ ("kafka" , "KafkaConsumer.__next__" ),
74
48
]
75
49
provider_name = "kafka"
76
50
name = "kafka"
77
51
78
- def _trace_send (
79
- self ,
80
- method ,
81
- topic ,
82
- value = None ,
83
- key = None ,
84
- headers = None ,
85
- partition = None ,
86
- timestamp_ms = None ,
87
- action = "send" ,
88
- ):
89
- span_name = "KafkaProducer#send to " + topic
90
- service = self .destination_info ["service" ]
91
- service ["resource" ] = service ["resource" ] + topic
92
- self .destination_info ["service" ] = service
52
+ def _trace_send (self , instance , wrapped , * args , destination_info = None , ** kwargs ):
53
+ topic = args [0 ] if args else kwargs ["topic" ]
54
+ headers = args [4 ] if len (args ) > 4 else kwargs .get ("headers" , None )
55
+
56
+ span_name = f"Kafka SEND to { topic } "
57
+ destination_info ["service" ]["resource" ] += topic
93
58
with capture_span (
94
59
name = span_name ,
95
60
span_type = "messaging" ,
96
61
span_subtype = self .provider_name ,
97
- span_action = action ,
62
+ span_action = "send" ,
98
63
extra = {
99
64
"message" : {"queue" : {"name" : topic }},
100
- "destination" : self . destination_info ,
65
+ "destination" : destination_info ,
101
66
},
102
67
) as span :
103
68
transaction = execution_context .get_transaction ()
104
69
if transaction :
105
- tp = transaction .trace_parent
106
- tp_string = tp .to_string ()
107
- # ID REPLACE SECTION START
108
- new_tp_string = tp_string .replace (transaction .id , span .id )
70
+ tp = transaction .trace_parent .copy_from (span_id = span .id )
109
71
if headers :
110
- headers .append (("trace" , bytes ( new_tp_string )))
72
+ headers .append ((constants . TRACEPARENT_BINARY_HEADER_NAME , tp . to_binary ( )))
111
73
else :
112
- headers = [("trace" , bytes (new_tp_string ))]
113
- # ID REPLACE SECTION STOP
114
- result = method (
115
- topic ,
116
- value = value ,
117
- key = key ,
118
- headers = headers ,
119
- partition = partition ,
120
- timestamp_ms = timestamp_ms ,
121
- )
74
+ headers = [(constants .TRACEPARENT_BINARY_HEADER_NAME , tp .to_binary ())]
75
+ if len (args ) > 4 :
76
+ args = list (args )
77
+ args [4 ] = headers
78
+ else :
79
+ kwargs ["headers" ] = headers
80
+ result = wrapped (* args , ** kwargs )
81
+ if instance and instance ._metadata .controller :
82
+ address = instance ._metadata .controller [1 ]
83
+ port = instance ._metadata .controller [2 ]
84
+ span .context ["destination" ]["address" ] = address
85
+ span .context ["destination" ]["port" ] = port
122
86
return result
123
87
124
88
def call_if_sampling (self , module , method , wrapped , instance , args , kwargs ):
@@ -129,76 +93,86 @@ def call_if_sampling(self, module, method, wrapped, instance, args, kwargs):
129
93
return self .call (module , method , wrapped , instance , args , kwargs )
130
94
131
95
def call (self , module , method , wrapped , instance , args , kwargs ):
132
- topic = None
96
+ client = get_client ()
133
97
destination_info = {
134
98
"service" : {"name" : "kafka" , "resource" : "kafka/" , "type" : "messaging" },
135
99
}
136
- self . destination_info = destination_info
100
+
137
101
if method == "KafkaProducer.send" :
138
- address = None
139
- port = None
140
- time_start = time .time ()
141
- while not instance ._metadata .controller :
142
- if time .time () - time_start > 1 :
143
- break
144
- continue
145
- if instance :
146
- if instance ._metadata .controller :
147
- address = instance ._metadata .controller [1 ]
148
- port = instance ._metadata .controller [2 ]
149
- self .destination_info ["port" ] = port
150
- self .destination_info ["address" ] = address
151
- topic = args [0 ].encode ("utf-8" )
102
+ topic = args [0 ] if args else kwargs ["topic" ]
103
+ if client .should_ignore_topic (topic ) or not execution_context .get_transaction ():
104
+ return wrapped (* args , ** kwargs )
105
+ return self ._trace_send (instance , wrapped , destination_info = destination_info , * args , ** kwargs )
106
+
107
+ elif method == "KafkaConsumer.poll" :
152
108
transaction = execution_context .get_transaction ()
153
109
if transaction :
154
- return self ._trace_send (wrapped , topic , ** kwargs )
110
+ with capture_span (
111
+ name = "Kafka POLL" ,
112
+ span_type = "messaging" ,
113
+ span_subtype = self .provider_name ,
114
+ span_action = "poll" ,
115
+ extra = {
116
+ "destination" : destination_info ,
117
+ },
118
+ ) as span :
119
+ if instance ._subscription .subscription :
120
+ span .name += " from " + ", " .join (sorted (instance ._subscription .subscription ))
121
+ results = wrapped (* args , ** kwargs )
122
+ return results
155
123
156
- if method == "KafkaConsumer.next " :
124
+ elif method == "KafkaConsumer.__next__ " :
157
125
transaction = execution_context .get_transaction ()
158
126
if transaction and transaction .transaction_type != "messaging" :
159
- action = "consume"
127
+ # somebody started a transaction outside of the consumer,
128
+ # so we capture it as a span, and record the causal trace as a link
160
129
with capture_span (
161
130
name = "consumer" ,
162
131
span_type = "messaging" ,
163
132
span_subtype = self .provider_name ,
164
- span_action = action ,
133
+ span_action = "receive" ,
165
134
extra = {
166
135
"message" : {"queue" : {"name" : "" }},
167
- "destination" : self . destination_info ,
136
+ "destination" : destination_info ,
168
137
},
169
138
) as span :
170
- result = wrapped (* args , ** kwargs )
139
+ try :
140
+ result = wrapped (* args , ** kwargs )
141
+ except StopIteration :
142
+ span .cancel ()
143
+ raise
171
144
topic = result [0 ]
172
- new_trace_id = get_trace_id (result )
173
- service = self .destination_info ["service" ]
174
- service ["resource" ] = service ["resource" ] + topic
145
+ trace_parent = self .get_traceparent_from_result (result )
146
+ if trace_parent :
147
+ span .add_link (trace_parent )
148
+ destination_info ["service" ]["resource" ] += topic
175
149
span .context ["message" ]["queue" ]["name" ] = topic
176
- span .context ["destination" ]["service" ] = service
177
- span .name = "KafkaConsumer#receive from " + topic
178
- transaction .trace_parent = TraceParent .from_string (new_trace_id )
150
+ span .name = "Kafka RECEIVE from " + topic
179
151
return result
180
152
else :
181
- client = get_client ()
182
- if transaction and transaction .transaction_type == "messaging" :
153
+ # No transaction running, or this is a transaction started by us,
154
+ # so let's end it and start the next,
155
+ # unless a StopIteration is raised, at which point we do nothing.
156
+ if transaction :
183
157
client .end_transaction ()
184
-
185
158
result = wrapped (* args , ** kwargs )
186
159
topic = result [0 ]
187
- new_trace_id = None
188
- new_trace_id = get_trace_id (result )
189
-
190
- client .begin_transaction ("messaging" , trace_parent = None )
191
- transaction = execution_context .get_transaction ()
160
+ trace_parent = self .get_traceparent_from_result (result )
161
+ transaction = client .begin_transaction ("messaging" , trace_parent = trace_parent )
192
162
if result .timestamp_type == 0 :
193
163
current_time_millis = int (round (time .time () * 1000 ))
194
164
age = current_time_millis - result .timestamp
195
165
transaction .context = {
196
- "message" : {"age" : {"ms" : age }, "queue" : {"name" : topic }}
166
+ "message" : {"age" : {"ms" : age }, "queue" : {"name" : topic }},
167
+ "service" : {"framework" : {"name" : "Kafka" }},
197
168
}
198
- if new_trace_id :
199
- transaction .trace_parent = TraceParent .from_string (new_trace_id )
200
- t_name = "Kafka record from " + topic
201
- elasticapm .set_transaction_name (t_name , override = True )
169
+ transaction_name = "Kafka RECEIVE from " + topic
170
+ elasticapm .set_transaction_name (transaction_name , override = True )
202
171
res = constants .OUTCOME .SUCCESS
203
172
elasticapm .set_transaction_result (res , override = False )
204
173
return result
174
+
175
+ def get_traceparent_from_result (self , result ) -> Optional [TraceParent ]:
176
+ for k , v in result .headers :
177
+ if k == constants .TRACEPARENT_BINARY_HEADER_NAME :
178
+ return TraceParent .from_binary (v )
0 commit comments