Adopt orjson (#76)

A good chunk of the compute done by the Edge Proxy is JSON marshalling
and unmarshalling, and having a more performant library has an impact.

Todoist's production data shows an approximate:
- Reduction in average CPU usage of 1%
- Reduction in p99 latency of 5%
- Increased throughput by ~8%

With recent improvements, an ECS instance with 4 vCPU is able to handle
3000 requests per minute without noticeable variation to p99 latency.

Author: goncalossilva

requirements.in

@@ -6,6 +6,7 @@ flagsmith-flag-engine
 python-decouple
 python-dotenv
 pydantic
+orjson
 # sse-stuff
 sse-starlette
 asyncio

requirements.txt

@@ -28,6 +28,8 @@ idna==3.4
     #   requests
 marshmallow==3.20.1
     # via -r requirements.in
+orjson==3.9.7
+    # via -r requirements.in
 packaging==23.1
     # via marshmallow
 pydantic==1.10.12

src/cache.py

@@ -1,6 +1,7 @@
 import logging
 from datetime import datetime
 
+import orjson
 import requests
 
 from .exceptions import FlagsmithUnknownKeyError
@@ -22,7 +23,7 @@ def fetch_document(self, server_side_key):
             url, headers={"X-Environment-Key": server_side_key}
         )
         response.raise_for_status()
-        return response.json()
+        return orjson.loads(response.text)
 
     def refresh(self):
         received_error = False
@@ -31,10 +32,10 @@ def refresh(self):
                 self._cache[key_pair.client_side_key] = self.fetch_document(
                     key_pair.server_side_key
                 )
-            except requests.exceptions.HTTPError:
+            except (requests.exceptions.HTTPError, orjson.JSONDecodeError):
                 received_error = True
                 logger.exception(
-                    f"Received non 200 response for {key_pair.client_side_key}"
+                    f"Failed to fetch document for {key_pair.client_side_key}"
                 )
         if not received_error:
             self.last_updated_at = datetime.now()

src/main.py

@@ -3,7 +3,7 @@
 
 from fastapi import FastAPI, Header
 from fastapi.middleware.cors import CORSMiddleware
-from fastapi.responses import JSONResponse
+from fastapi.responses import ORJSONResponse
 from flag_engine.engine import (
     get_environment_feature_state,
     get_environment_feature_states,
@@ -33,7 +33,7 @@
 
 @app.exception_handler(FlagsmithUnknownKeyError)
 async def unknown_key_error(request, exc):
-    return JSONResponse(
+    return ORJSONResponse(
         status_code=401,
         content={
             "status": "unauthorized",
@@ -42,19 +42,19 @@ async def unknown_key_error(request, exc):
     )
 
 
-@app.get("/health", deprecated=True)
-@app.get("/proxy/health")
+@app.get("/health", response_class=ORJSONResponse, deprecated=True)
+@app.get("/proxy/health", response_class=ORJSONResponse)
 def health_check():
     with suppress(TypeError):
         last_updated = datetime.now() - cache_service.last_updated_at
         buffer = 30 * len(settings.environment_key_pairs)  # 30s per environment
         if last_updated.total_seconds() <= settings.api_poll_frequency + buffer:
-            return JSONResponse(status_code=200, content={"status": "ok"})
+            return ORJSONResponse(status_code=200, content={"status": "ok"})
 
-    return JSONResponse(status_code=500, content={"status": "error"})
+    return ORJSONResponse(status_code=500, content={"status": "error"})
 
 
-@app.get("/api/v1/flags/")
+@app.get("/api/v1/flags/", response_class=ORJSONResponse)
 def flags(feature: str = None, x_environment_key: str = Header(None)):
     environment_document = cache_service.get_environment(x_environment_key)
     environment = build_environment_model(environment_document)
@@ -66,7 +66,7 @@ def flags(feature: str = None, x_environment_key: str = Header(None)):
             feature_states=[feature_state],
             environment=environment,
         ):
-            return JSONResponse(
+            return ORJSONResponse(
                 status_code=404,
                 content={
                     "status": "not_found",
@@ -83,10 +83,10 @@ def flags(feature: str = None, x_environment_key: str = Header(None)):
         )
         data = map_feature_states_to_response_data(feature_states)
 
-    return JSONResponse(data)
+    return ORJSONResponse(data)
 
 
-@app.post("/api/v1/identities/")
+@app.post("/api/v1/identities/", response_class=ORJSONResponse)
 def identity(
     input_data: IdentityWithTraits,
     x_environment_key: str = Header(None),
@@ -112,7 +112,7 @@ def identity(
             identity_hash_key=identity.composite_key,
         ),
     }
-    return JSONResponse(data)
+    return ORJSONResponse(data)
 
 
 @app.on_event("startup")

tests/test_cache.py

@@ -1,3 +1,5 @@
+import unittest.mock
+
 import pytest
 import requests
 
@@ -16,14 +18,19 @@
 
 def test_refresh_makes_correct_http_call(mocker):
     # Given
-    mocked_session = mocker.patch("src.cache.requests.Session")
+    mocked_get = mocker.patch("src.cache.requests.Session.get")
+    mocked_get.side_effect = [
+        unittest.mock.AsyncMock(text='{"key1": "value1"}'),
+        unittest.mock.AsyncMock(text='{"key2": "value2"}'),
+    ]
     mocked_datetime = mocker.patch("src.cache.datetime")
     cache_service = CacheService(settings)
 
     # When
     cache_service.refresh()
+
     # Then
-    mocked_session.return_value.get.assert_has_calls(
+    mocked_get.assert_has_calls(
         [
             mocker.call(
                 f"{settings.api_url}/environment-document/",