Use orjson as dramatiq encoder

Signed-off-by: Federico Busetti <[email protected]>

Author: febus982

poetry.lock

@@ -34,6 +34,7 @@ opentelemetry-instrumentation-celery = "*"
 opentelemetry-instrumentation-httpx = "*"
 opentelemetry-instrumentation-sqlalchemy = "*"
 opentelemetry-instrumentor-dramatiq = "*"
+orjson = "^3.10.12"
 pydantic = "^2.2.1"
 pydantic-settings = "^2.0.3"
 python = ">=3.9,<3.14"

pyproject.toml

@@ -1,12 +1,27 @@
-from dramatiq import set_broker
+from dramatiq import set_broker, set_encoder
 from dramatiq.brokers.stub import StubBroker
 from dramatiq.brokers.redis import RedisBroker
+from dramatiq.encoder import Encoder, DecodeError, MessageData
 from dramatiq.middleware import AsyncIO
 from opentelemetry_instrumentor_dramatiq import DramatiqInstrumentor
-
+import orjson
 from .config import AppConfig
 
 
+class ORJSONEncoder(Encoder):
+    """Encodes messages as JSON. orjson is much faster than json.
+    """
+
+    def encode(self, data: MessageData) -> bytes:
+        return orjson.dumps(data)
+
+    def decode(self, data: bytes) -> MessageData:
+        try:
+            return orjson.loads(data)
+        except orjson.JSONDecodeError as e:
+            raise DecodeError("failed to decode message %r" % (data,), data, e) from None
+
+
 def init_dramatiq(config: AppConfig):
     DramatiqInstrumentor().instrument()
     if config.ENVIRONMENT == "test":
@@ -18,3 +33,4 @@ def init_dramatiq(config: AppConfig):
         raise RuntimeError("Running a non-test environment without Redis URL set")
     broker.add_middleware(AsyncIO())
     set_broker(broker)
+    set_encoder(ORJSONEncoder())