From 1adeee308fb892131d6401108b451d0c020be5c2 Mon Sep 17 00:00:00 2001
From: Thomas Sileo <t@a4.io>
Date: Wed, 11 Jul 2018 20:04:48 +0200
Subject: [PATCH] Cache actors

---
 activitypub.py   | 30 +++++++++++++++++++++++++++++-
 requirements.txt |  1 +
 2 files changed, 30 insertions(+), 1 deletion(-)

diff --git a/activitypub.py b/activitypub.py
index 336ede5..1a6a9fe 100644
--- a/activitypub.py
+++ b/activitypub.py
@@ -10,6 +10,7 @@ from typing import Optional
 from bson.objectid import ObjectId
 from feedgen.feed import FeedGenerator
 from html2text import html2text
+from cachetools import LRUCache
 
 import tasks
 from config import BASE_URL
@@ -31,6 +32,9 @@ from utils.media import Kind
 logger = logging.getLogger(__name__)
 
 
+ACTORS_CACHE = LRUCache(maxsize=256)
+
+
 def _remove_id(doc: ap.ObjectType) -> ap.ObjectType:
     """Helper for removing MongoDB's `_id` field."""
     doc = doc.copy()
@@ -140,7 +144,7 @@ class MicroblogPubBackend(Backend):
             )
         )
 
-    def fetch_iri(self, iri: str) -> ap.ObjectType:
+    def _fetch_iri(self, iri: str) -> ap.ObjectType:
         if iri == ME["id"]:
             return ME
 
@@ -168,6 +172,30 @@ class MicroblogPubBackend(Backend):
         # Fetch the URL via HTTP
         return super().fetch_iri(iri)
 
+    def fetch_iri(self, iri: str) -> ap.ObjectType:
+        if iri == ME["id"]:
+            return ME
+
+        if iri in ACTORS_CACHE:
+            return ACTORS_CACHE[iri]
+
+        data = DB.actors.find_one({"remote_id": iri})
+        if data:
+            ACTORS_CACHE[iri] = data['data']
+            return data['data']
+
+        data = self._fetch_iri(iri)
+        if ap._has_type(data["type"], ap.ACTOR_TYPES):
+            # Cache the actor
+            DB.actors.update_one(
+                {"remote_id": iri},
+                {"$set": {"remote_id": iri, "data": data}},
+                upsert=True,
+            )
+            ACTORS_CACHE[iri] = data
+
+        return data
+
     @ensure_it_is_me
     def inbox_check_duplicate(self, as_actor: ap.Person, iri: str) -> bool:
         return bool(DB.activities.find_one({"box": Box.INBOX.value, "remote_id": iri}))
diff --git a/requirements.txt b/requirements.txt
index 21ce34c..d2eec2d 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -20,3 +20,4 @@ git+https://github.com/erikriver/opengraph.git
 git+https://github.com/tsileo/little-boxes.git
 pyyaml
 pillow
+cachetools