From 037d29b3daa82dcb284fdda3f501e3d19f58b4ab Mon Sep 17 00:00:00 2001
From: Justin Zhang <judtin@seas.upenn.edu>
Date: Thu, 25 Apr 2024 15:05:57 -0400
Subject: [PATCH] In-memory caching JWT decode/serialization

---
 analytics/analytics.py            | 30 +++++++++++++++++++++++-------
 tests/analytics/test_analytics.py | 13 ++++++++-----
 tests/settings.py                 |  2 ++
 tox.ini                           |  4 ++--
 4 files changed, 35 insertions(+), 14 deletions(-)

diff --git a/analytics/analytics.py b/analytics/analytics.py
index e2448ae..dffec6c 100644
--- a/analytics/analytics.py
+++ b/analytics/analytics.py
@@ -1,4 +1,5 @@
 import json
+import time
 from concurrent.futures import ThreadPoolExecutor
 from enum import IntEnum
 from typing import Optional
@@ -65,17 +66,32 @@ def __new__(cls, *args, **kwargs):
 
     def __init__(self): 
         self.executor = ThreadPoolExecutor(max_workers=self.POOL_SIZE)
+        self.session = NoRebuildAuthSession()
 
-    def submit(self, txn: AnalyticsTxn):
-        _refresh_if_outdated()
+        self.expires_at = None
+        self.headers = dict()
+        
+        # Local caching of expiration date and headers
+        self._refresh_expires_at()
+        self._refresh_headers()
+
+    def _refresh_expires_at(self):
+        self.expires_at = json.loads(container.access_jwt.claims)["exp"]
 
-        headers = {
+    def _refresh_headers(self):
+        self.headers = {
             "Authorization": f"Bearer {container.access_jwt.serialize()}",
             "Content-Type": "application/json"
         }
-        self.executor.submit(self.request_job, txn.to_json(), headers)
 
-    def request_job(self, json, headers):
-        session = NoRebuildAuthSession()
-        session.post(url=self.ANALYTICS_URL, json=json, headers=headers)
+    def submit(self, txn: AnalyticsTxn):
+        # Offer a 30 second buffer to refresh
+        if time.time() < self.expires_at - 30: 
+            _refresh_if_outdated()
+            self._refresh_expires_at()
+            self._refresh_headers()
+
+        self.executor.submit(self.send_message, txn.to_json())
 
+    def send_message(self, json):
+        self.session.post(url=self.ANALYTICS_URL, json=json, headers=self.headers)
diff --git a/tests/analytics/test_analytics.py b/tests/analytics/test_analytics.py
index a19944f..7ad31c9 100644
--- a/tests/analytics/test_analytics.py
+++ b/tests/analytics/test_analytics.py
@@ -16,7 +16,6 @@ def test_analytics_txn(self):
 
         txn = AnalyticsTxn(**data)
         data_json = txn.to_json()
-        # data_dict = json.dumps(data_json)
 
         self.assertEqual(Product.MOBILE_BACKEND.value, int(data_json["product"]))
         self.assertIsNone(data_json["pennkey"])
@@ -33,11 +32,15 @@ def setUp(self):
     def test_submit(self):
         data = {
             "product": Product.MOBILE_BACKEND,
-            "pennkey": "hi",
-            "data": [{"key": "backend", "value": "data"}],
+            "pennkey": "judtin",
+            "data": [{"key": "backend", "value": "some data"}],
         }
 
-        txn = AnalyticsTxn(**data)
+        for _ in range(20):
+            data["product"] = Product((data["product"].value + 1) % len(Product))
+            txn = AnalyticsTxn(**data)
+            self.analytics_wrapper.submit(txn)
+
+        self.analytics_wrapper.executor.shutdown(wait=True)
 
-        self.analytics_wrapper.submit(txn)
         assert False
\ No newline at end of file
diff --git a/tests/settings.py b/tests/settings.py
index c9892b3..d16374d 100644
--- a/tests/settings.py
+++ b/tests/settings.py
@@ -10,6 +10,8 @@
     "django.contrib.admin",
     "django.contrib.messages",
     "accounts.apps.AccountsConfig",
+    "identity.apps.IdentityConfig",
+    "analytics.apps.AnalyticsConfig",
     "tests",
 )
 
diff --git a/tox.ini b/tox.ini
index e417f76..197f70c 100644
--- a/tox.ini
+++ b/tox.ini
@@ -1,9 +1,9 @@
 [tox]
 isolated_build = true
 envlist =
-    ; lint,
+    lint,
     py311-django{502},
-    ; sentry-django{30,31},
+    sentry-django{502},
 
 [testenv]
 allowlist_externals = poetry