Skip to content

Commit 2f4c293

Browse files
jan-auerjan-auer
authored
ref(spans): Use orjson across the entire buffer (#94162)
Switches to `orjson` like the rest of the codebase.
1 parent db0c633 commit 2f4c293

File tree

5 files changed

+42
-42
lines changed

5 files changed

+42
-42
lines changed

‎src/sentry/spans/buffer.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -69,7 +69,7 @@
6969
from collections.abc import Generator, MutableMapping, Sequence
7070
from typing import Any, NamedTuple
7171

72-
import rapidjson
72+
import orjson
7373
import zstandard
7474
from django.conf import settings
7575
from django.utils.functional import cached_property
@@ -418,7 +418,7 @@ def flush_segments(self, now: int) -> dict[SegmentKey, FlushedSegment]:
418418
has_root_span = False
419419
metrics.timing("spans.buffer.flush_segments.num_spans_per_segment", len(segment))
420420
for payload in segment:
421-
val = rapidjson.loads(payload)
421+
val = orjson.loads(payload)
422422
old_segment_id = val.get("segment_id")
423423
outcome = "same" if old_segment_id == segment_span_id else "different"
424424

‎src/sentry/spans/consumers/process/factory.py

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -4,7 +4,7 @@
44
from functools import partial
55
from typing import cast
66

7-
import rapidjson
7+
import orjson
88
import sentry_sdk
99
from arroyo.backends.kafka.consumer import KafkaPayload
1010
from arroyo.dlq import InvalidMessage
@@ -143,7 +143,7 @@ def process_batch(
143143
min_timestamp = timestamp
144144

145145
with metrics.timer("spans.buffer.process_batch.decode"):
146-
val = cast(SpanEvent, rapidjson.loads(payload.value))
146+
val = cast(SpanEvent, orjson.loads(payload.value))
147147

148148
if killswitches.killswitch_matches_context(
149149
"spans.drop-in-buffer",

‎tests/sentry/spans/consumers/process/test_consumer.py

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,8 @@
11
import time
22
from datetime import datetime
33

4+
import orjson
45
import pytest
5-
import rapidjson
66
from arroyo.backends.kafka import KafkaPayload
77
from arroyo.types import BrokerValue, Message, Partition, Topic
88

@@ -40,14 +40,14 @@ def add_commit(offsets, force=False):
4040
offset=1,
4141
payload=KafkaPayload(
4242
None,
43-
rapidjson.dumps(
43+
orjson.dumps(
4444
{
4545
"project_id": 12,
4646
"span_id": "a" * 16,
4747
"trace_id": "b" * 32,
4848
"end_timestamp_precise": 1700000000.0,
4949
}
50-
).encode("ascii"),
50+
),
5151
[],
5252
),
5353
timestamp=datetime.now(),
@@ -66,7 +66,7 @@ def add_commit(offsets, force=False):
6666

6767
(msg,) = messages
6868

69-
assert rapidjson.loads(msg.value) == {
69+
assert orjson.loads(msg.value) == {
7070
"spans": [
7171
{
7272
"data": {

‎tests/sentry/spans/consumers/process/test_flusher.py

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11
import time
22
from time import sleep
33

4-
import rapidjson
4+
import orjson
55
from arroyo.processing.strategies.noop import Noop
66

77
from sentry.spans.buffer import Span, SpansBuffer
@@ -10,8 +10,8 @@
1010
from tests.sentry.spans.test_buffer import DEFAULT_OPTIONS
1111

1212

13-
def _payload(span_id: bytes) -> bytes:
14-
return rapidjson.dumps({"span_id": span_id}).encode("ascii")
13+
def _payload(span_id: str) -> bytes:
14+
return orjson.dumps({"span_id": span_id})
1515

1616

1717
@override_options({**DEFAULT_OPTIONS, "spans.buffer.max-flush-segments": 1})
@@ -39,31 +39,31 @@ def append(msg):
3939

4040
spans = [
4141
Span(
42-
payload=_payload(b"a" * 16),
42+
payload=_payload("a" * 16),
4343
trace_id=trace_id,
4444
span_id="a" * 16,
4545
parent_span_id="b" * 16,
4646
project_id=1,
4747
end_timestamp_precise=now,
4848
),
4949
Span(
50-
payload=_payload(b"d" * 16),
50+
payload=_payload("d" * 16),
5151
trace_id=trace_id,
5252
span_id="d" * 16,
5353
parent_span_id="b" * 16,
5454
project_id=1,
5555
end_timestamp_precise=now,
5656
),
5757
Span(
58-
payload=_payload(b"c" * 16),
58+
payload=_payload("c" * 16),
5959
trace_id=trace_id,
6060
span_id="c" * 16,
6161
parent_span_id="b" * 16,
6262
project_id=1,
6363
end_timestamp_precise=now,
6464
),
6565
Span(
66-
payload=_payload(b"b" * 16),
66+
payload=_payload("b" * 16),
6767
trace_id=trace_id,
6868
span_id="b" * 16,
6969
parent_span_id=None,

‎tests/sentry/spans/test_buffer.py

Lines changed: 27 additions & 27 deletions
Original file line numberDiff line numberDiff line change
@@ -4,8 +4,8 @@
44
from collections.abc import Sequence
55
from unittest import mock
66

7+
import orjson
78
import pytest
8-
import rapidjson
99
from sentry_redis_tools.clients import StrictRedis
1010

1111
from sentry.spans.buffer import FlushedSegment, OutputSpan, SegmentKey, Span, SpansBuffer
@@ -38,8 +38,8 @@ def _segment_id(project_id: int, trace_id: str, span_id: str) -> SegmentKey:
3838
return f"span-buf:z:{{{project_id}:{trace_id}}}:{span_id}".encode("ascii")
3939

4040

41-
def _payload(span_id: bytes) -> bytes:
42-
return rapidjson.dumps({"span_id": span_id}).encode("ascii")
41+
def _payload(span_id: str) -> bytes:
42+
return orjson.dumps({"span_id": span_id})
4343

4444

4545
def _output_segment(span_id: bytes, segment_id: bytes, is_segment: bool) -> OutputSpan:
@@ -134,31 +134,31 @@ def process_spans(spans: Sequence[Span | _SplitBatch], buffer: SpansBuffer, now)
134134
itertools.permutations(
135135
[
136136
Span(
137-
payload=_payload(b"a" * 16),
137+
payload=_payload("a" * 16),
138138
trace_id="a" * 32,
139139
span_id="a" * 16,
140140
parent_span_id="b" * 16,
141141
project_id=1,
142142
end_timestamp_precise=1700000000.0,
143143
),
144144
Span(
145-
payload=_payload(b"d" * 16),
145+
payload=_payload("d" * 16),
146146
trace_id="a" * 32,
147147
span_id="d" * 16,
148148
parent_span_id="b" * 16,
149149
project_id=1,
150150
end_timestamp_precise=1700000000.0,
151151
),
152152
Span(
153-
payload=_payload(b"c" * 16),
153+
payload=_payload("c" * 16),
154154
trace_id="a" * 32,
155155
span_id="c" * 16,
156156
parent_span_id="b" * 16,
157157
project_id=1,
158158
end_timestamp_precise=1700000000.0,
159159
),
160160
Span(
161-
payload=_payload(b"b" * 16),
161+
payload=_payload("b" * 16),
162162
trace_id="a" * 32,
163163
span_id="b" * 16,
164164
parent_span_id=None,
@@ -203,7 +203,7 @@ def test_basic(buffer: SpansBuffer, spans):
203203
itertools.permutations(
204204
[
205205
Span(
206-
payload=_payload(b"d" * 16),
206+
payload=_payload("d" * 16),
207207
trace_id="a" * 32,
208208
span_id="d" * 16,
209209
parent_span_id="b" * 16,
@@ -212,15 +212,15 @@ def test_basic(buffer: SpansBuffer, spans):
212212
),
213213
_SplitBatch(),
214214
Span(
215-
payload=_payload(b"b" * 16),
215+
payload=_payload("b" * 16),
216216
trace_id="a" * 32,
217217
span_id="b" * 16,
218218
parent_span_id="a" * 16,
219219
project_id=1,
220220
end_timestamp_precise=1700000000.0,
221221
),
222222
Span(
223-
payload=_payload(b"a" * 16),
223+
payload=_payload("a" * 16),
224224
trace_id="a" * 32,
225225
span_id="a" * 16,
226226
parent_span_id=None,
@@ -229,7 +229,7 @@ def test_basic(buffer: SpansBuffer, spans):
229229
end_timestamp_precise=1700000000.0,
230230
),
231231
Span(
232-
payload=_payload(b"c" * 16),
232+
payload=_payload("c" * 16),
233233
trace_id="a" * 32,
234234
span_id="c" * 16,
235235
parent_span_id="a" * 16,
@@ -273,39 +273,39 @@ def test_deep(buffer: SpansBuffer, spans):
273273
itertools.permutations(
274274
[
275275
Span(
276-
payload=_payload(b"e" * 16),
276+
payload=_payload("e" * 16),
277277
trace_id="a" * 32,
278278
span_id="e" * 16,
279279
parent_span_id="d" * 16,
280280
project_id=1,
281281
end_timestamp_precise=1700000000.0,
282282
),
283283
Span(
284-
payload=_payload(b"d" * 16),
284+
payload=_payload("d" * 16),
285285
trace_id="a" * 32,
286286
span_id="d" * 16,
287287
parent_span_id="b" * 16,
288288
project_id=1,
289289
end_timestamp_precise=1700000000.0,
290290
),
291291
Span(
292-
payload=_payload(b"b" * 16),
292+
payload=_payload("b" * 16),
293293
trace_id="a" * 32,
294294
span_id="b" * 16,
295295
parent_span_id="c" * 16,
296296
project_id=1,
297297
end_timestamp_precise=1700000000.0,
298298
),
299299
Span(
300-
payload=_payload(b"c" * 16),
300+
payload=_payload("c" * 16),
301301
trace_id="a" * 32,
302302
span_id="c" * 16,
303303
parent_span_id="a" * 16,
304304
project_id=1,
305305
end_timestamp_precise=1700000000.0,
306306
),
307307
Span(
308-
payload=_payload(b"a" * 16),
308+
payload=_payload("a" * 16),
309309
trace_id="a" * 32,
310310
span_id="a" * 16,
311311
parent_span_id=None,
@@ -351,31 +351,31 @@ def test_deep2(buffer: SpansBuffer, spans):
351351
itertools.permutations(
352352
[
353353
Span(
354-
payload=_payload(b"c" * 16),
354+
payload=_payload("c" * 16),
355355
trace_id="a" * 32,
356356
span_id="c" * 16,
357357
parent_span_id="b" * 16,
358358
project_id=1,
359359
end_timestamp_precise=1700000000.0,
360360
),
361361
Span(
362-
payload=_payload(b"d" * 16),
362+
payload=_payload("d" * 16),
363363
trace_id="a" * 32,
364364
span_id="d" * 16,
365365
parent_span_id="b" * 16,
366366
project_id=1,
367367
end_timestamp_precise=1700000000.0,
368368
),
369369
Span(
370-
payload=_payload(b"e" * 16),
370+
payload=_payload("e" * 16),
371371
trace_id="a" * 32,
372372
span_id="e" * 16,
373373
parent_span_id="b" * 16,
374374
project_id=1,
375375
end_timestamp_precise=1700000000.0,
376376
),
377377
Span(
378-
payload=_payload(b"b" * 16),
378+
payload=_payload("b" * 16),
379379
trace_id="a" * 32,
380380
span_id="b" * 16,
381381
parent_span_id=None,
@@ -427,7 +427,7 @@ def test_parent_in_other_project(buffer: SpansBuffer, spans):
427427
shallow_permutations(
428428
[
429429
Span(
430-
payload=_payload(b"c" * 16),
430+
payload=_payload("c" * 16),
431431
trace_id="a" * 32,
432432
span_id="c" * 16,
433433
parent_span_id="d" * 16,
@@ -436,23 +436,23 @@ def test_parent_in_other_project(buffer: SpansBuffer, spans):
436436
end_timestamp_precise=1700000000.0,
437437
),
438438
Span(
439-
payload=_payload(b"d" * 16),
439+
payload=_payload("d" * 16),
440440
trace_id="a" * 32,
441441
span_id="d" * 16,
442442
parent_span_id="b" * 16,
443443
project_id=1,
444444
end_timestamp_precise=1700000000.0,
445445
),
446446
Span(
447-
payload=_payload(b"e" * 16),
447+
payload=_payload("e" * 16),
448448
trace_id="a" * 32,
449449
span_id="e" * 16,
450450
parent_span_id="b" * 16,
451451
project_id=1,
452452
end_timestamp_precise=1700000000.0,
453453
),
454454
Span(
455-
payload=_payload(b"b" * 16),
455+
payload=_payload("b" * 16),
456456
trace_id="a" * 32,
457457
span_id="b" * 16,
458458
parent_span_id=None,
@@ -507,7 +507,7 @@ def test_parent_in_other_project_and_nested_is_segment_span(buffer: SpansBuffer,
507507
def test_flush_rebalance(buffer: SpansBuffer):
508508
spans = [
509509
Span(
510-
payload=_payload(b"a" * 16),
510+
payload=_payload("a" * 16),
511511
trace_id="a" * 32,
512512
span_id="a" * 16,
513513
parent_span_id=None,
@@ -545,14 +545,14 @@ def test_compression_functionality(compression_level):
545545
buffer = SpansBuffer(assigned_shards=list(range(32)))
546546

547547
def make_payload(span_id: str):
548-
return rapidjson.dumps(
548+
return orjson.dumps(
549549
{
550550
"span_id": span_id,
551551
"trace_id": "a" * 32,
552552
"data": {"message": "x" * 1000},
553553
"extra_data": {"field": "y" * 500},
554554
}
555-
).encode("ascii")
555+
)
556556

557557
spans = [
558558
Span(

0 commit comments

Comments
 (0)