From a646c9656e34032a532beec533ba2d83bb1e55c9 Mon Sep 17 00:00:00 2001
From: Your Name <you@example.com>
Date: Sat, 8 Jul 2023 17:30:56 -0400
Subject: [PATCH] rate limit fetch

---
 catalog/api.py                       |  5 +++--
 catalog/search/models.py             | 12 ++++++++++++
 catalog/search/views.py              |  8 +++++---
 catalog/templates/fetch_pending.html | 16 ++++++++++------
 users/models.py                      |  4 ++--
 5 files changed, 32 insertions(+), 13 deletions(-)

diff --git a/catalog/api.py b/catalog/api.py
index 240dd107..89a1dbc3 100644
--- a/catalog/api.py
+++ b/catalog/api.py
@@ -6,7 +6,7 @@ from common.api import *
 from .models import *
 from .common import *
 from .sites import *
-from .search.models import enqueue_fetch, query_index
+from .search.models import enqueue_fetch, query_index, get_fetch_lock
 
 
 class SearchResult(Schema):
@@ -70,7 +70,8 @@ def fetch_item(request, url: str):
     item = site.get_item()
     if item:
         return 200, item
-    enqueue_fetch(url, False)
+    if get_fetch_lock():
+        enqueue_fetch(url, False)
     return 202, {"message": "Fetch in progress"}
 
 
diff --git a/catalog/search/models.py b/catalog/search/models.py
index e9fbdf91..210b347a 100644
--- a/catalog/search/models.py
+++ b/catalog/search/models.py
@@ -9,6 +9,7 @@ from django.core.cache import cache
 import hashlib
 from .typesense import Indexer as TypeSenseIndexer
 from auditlog.context import set_actor
+from django.core.cache import cache
 
 # from .meilisearch import Indexer as MeiliSearchIndexer
 
@@ -98,6 +99,17 @@ def query_index(keywords, category=None, tag=None, page=1, prepare_external=True
     return items, result.num_pages, result.count, duplicated_items
 
 
+_fetch_lock_key = "_fetch_lock"
+_fetch_lock_ttl = 2
+
+
+def get_fetch_lock():
+    if cache.get(_fetch_lock_key):
+        return False
+    cache.set(_fetch_lock_key, 1, timeout=_fetch_lock_ttl)
+    return True
+
+
 def enqueue_fetch(url, is_refetch, user=None):
     job_id = "fetch_" + hashlib.md5(url.encode()).hexdigest()
     in_progress = False
diff --git a/catalog/search/views.py b/catalog/search/views.py
index 36feaf97..b88cebc5 100644
--- a/catalog/search/views.py
+++ b/catalog/search/views.py
@@ -16,7 +16,7 @@ from rq.job import Job
 from .external import ExternalSources
 from django.core.cache import cache
 import hashlib
-from .models import query_index, enqueue_fetch
+from .models import get_fetch_lock, query_index, enqueue_fetch
 
 _logger = logging.getLogger(__name__)
 
@@ -68,7 +68,9 @@ def fetch(request, url, is_refetch: bool = False, site: AbstractSite | None = No
                 "!refetch": [url, None],
             }
         )
-    job_id = enqueue_fetch(url, is_refetch, request.user)
+    job_id = None
+    if is_refetch or get_fetch_lock():
+        job_id = enqueue_fetch(url, is_refetch, request.user)
     return render(
         request,
         "fetch_pending.html",
@@ -98,7 +100,7 @@ def search(request):
             },
         )
 
-    if request.user.is_authenticated and keywords.find("://") > 0:
+    if keywords.find("://") > 0:
         site = SiteManager.get_site_by_url(keywords)
         if site:
             return fetch(request, keywords, False, site)
diff --git a/catalog/templates/fetch_pending.html b/catalog/templates/fetch_pending.html
index 1cd2ef56..aab03d11 100644
--- a/catalog/templates/fetch_pending.html
+++ b/catalog/templates/fetch_pending.html
@@ -21,12 +21,16 @@
     <main>
       <div class="grid__main">
         <article>
-          <h5>正在从{{ site.SITE_NAME.label }}获取</h5>
-          <div hx-get="{% url 'catalog:fetch_refresh' job_id %}"
-               hx-trigger="load delay:2s"
-               hx-swap="outerHTML">
-            <i class="fa-solid fa-compact-disc fa-spin loading"></i>
-          </div>
+          {% if job_id %}
+            <h5>正在从{{ site.SITE_NAME.label }}获取</h5>
+            <div hx-get="{% url 'catalog:fetch_refresh' job_id %}"
+                 hx-trigger="load delay:2s"
+                 hx-swap="outerHTML">
+              <i class="fa-solid fa-compact-disc fa-spin loading"></i>
+            </div>
+          {% else %}
+            <h5>获取系统繁忙，请稍等几秒钟再搜索</h5>
+          {% endif %}
         </article>
       </div>
       <aside class="grid__aside bottom">
diff --git a/users/models.py b/users/models.py
index 41effa58..151762f0 100644
--- a/users/models.py
+++ b/users/models.py
@@ -18,7 +18,7 @@ from django.db.models import Q, F, Value
 from django.db.models.functions import Concat
 from django.templatetags.static import static
 import hashlib
-
+from loguru import logger
 
 RESERVED_USERNAMES = [
     "connect",
@@ -394,7 +394,7 @@ class User(AbstractUser):
             self.merge_relationships()
             updated = True
         elif code == 401:
-            print(f"401 {self}")
+            logger.error(f"401 {self}")
             self.mastodon_token = ""
         return updated