更换cache使用方式

conf/views.py

@@ -1,7 +1,6 @@
 import hashlib
 
 from django.utils import timezone
-from django_redis import get_redis_connection
 
 from account.decorators import super_admin_required
 from judge.languages import languages, spj_languages
@@ -129,8 +128,7 @@ class JudgeServerHeartbeatAPI(CSRFExemptAPIView):
                                        last_heartbeat=timezone.now(),
                                        )
             # 新server上线 处理队列中的，防止没有新的提交而导致一直waiting
-            conn = get_redis_connection("JudgeQueue")
-            process_pending_task(conn)
+            process_pending_task()
 
         return self.success()
 

judge/dispatcher.py

@@ -7,7 +7,6 @@ import requests
 from django.core.cache import cache
 from django.db import transaction
 from django.db.models import F
-from django_redis import get_redis_connection
 
 from account.models import User
 from conf.models import JudgeServer, JudgeServerToken
@@ -15,18 +14,18 @@ from contest.models import ContestRuleType, ACMContestRank, OIContestRank
 from judge.languages import languages
 from problem.models import Problem, ProblemRuleType, ContestProblem
 from submission.models import JudgeStatus, Submission
+from utils.cache import judge_queue_cache
+from utils.constants import CacheKey
 
 logger = logging.getLogger(__name__)
 
-WAITING_QUEUE = "waiting_queue"
-
 
 # 继续处理在队列中的问题
-def process_pending_task(redis_conn):
-    if redis_conn.llen(WAITING_QUEUE):
+def process_pending_task():
+    if judge_queue_cache.llen(CacheKey.waiting_queue):
         # 防止循环引入
         from judge.tasks import judge_task
-        data = json.loads(redis_conn.rpop(WAITING_QUEUE))
+        data = json.loads(judge_queue_cache.rpop(CacheKey.waiting_queue))
         judge_task.delay(**data)
 
 
@@ -34,7 +33,7 @@ class JudgeDispatcher(object):
     def __init__(self, submission_id, problem_id):
         token = JudgeServerToken.objects.first().token
         self.token = hashlib.sha256(token.encode("utf-8")).hexdigest()
-        self.redis_conn = get_redis_connection("JudgeQueue")
+        self.redis_conn = judge_queue_cache
         self.submission = Submission.objects.get(pk=submission_id)
         if self.submission.contest_id:
             self.problem = ContestProblem.objects.select_related("contest")\
@@ -77,7 +76,7 @@ class JudgeDispatcher(object):
         server = self.choose_judge_server()
         if not server:
             data = {"submission_id": self.submission.id, "problem_id": self.problem.id}
-            self.redis_conn.lpush(WAITING_QUEUE, json.dumps(data))
+            self.redis_conn.lpush(CacheKey.waiting_queue, json.dumps(data))
             return
 
         sub_config = list(filter(lambda item: self.submission.language == item["name"], languages))[0]
@@ -130,7 +129,7 @@ class JudgeDispatcher(object):
         else:
             self.update_problem_status()
         # 至此判题结束，尝试处理任务队列中剩余的任务
-        process_pending_task(self.redis_conn)
+        process_pending_task()
 
     def compile_spj(self, service_url, src, spj_version, spj_compile_config, test_case_id):
         data = {"src": src, "spj_version": spj_version,

oj/settings.py

@@ -159,6 +159,42 @@ REST_FRAMEWORK = {
     )
 }
 
+CACHE_JUDGE_QUEUE = "judge_queue"
+CACHE_THROTTLING = "throttling"
+
+
+CACHES = {
+    "default": {
+        "BACKEND": "django_redis.cache.RedisCache",
+        "LOCATION": "redis://127.0.0.1:6379/1",
+        "OPTIONS": {
+            "CLIENT_CLASS": "django_redis.client.DefaultClient",
+        }
+    },
+    CACHE_JUDGE_QUEUE: {
+        "BACKEND": "django_redis.cache.RedisCache",
+        "LOCATION": "redis://127.0.0.1:6379/2",
+        "OPTIONS": {
+            "CLIENT_CLASS": "django_redis.client.DefaultClient",
+        }
+    },
+    CACHE_THROTTLING: {
+        "BACKEND": "django_redis.cache.RedisCache",
+        "LOCATION": "redis://127.0.0.1:6379/3",
+        "OPTIONS": {
+            "CLIENT_CLASS": "django_redis.client.DefaultClient",
+        }
+    }
+}
+
+# For celery
+REDIS_QUEUE = {
+    "host": "127.0.0.1",
+    "port": 6379,
+    "db": 4
+}
+
+
 # for celery
 BROKER_URL = 'redis://%s:%s/%s' % (REDIS_QUEUE["host"], str(REDIS_QUEUE["port"]), str(REDIS_QUEUE["db"]))
 CELERY_ACCEPT_CONTENT = ["json"]

submission/urls/oj.py

@@ -1,9 +1,8 @@
 from django.conf.urls import url
 
-from ..views.oj import SubmissionAPI, SubmissionListAPI, ContestSubmissionListAPI
+from ..views.oj import SubmissionAPI, SubmissionListAPI
 
 urlpatterns = [
     url(r"^submission/?$", SubmissionAPI.as_view(), name="submission_api"),
-    url(r"^submissions/?$", SubmissionListAPI.as_view(), name="submission_list_api"),
-    url(r"^contest/submissions/?$", ContestSubmissionListAPI.as_view(), name="contest_submission_list_api"),
+    url(r"^submissions/?$", SubmissionListAPI.as_view(), name="submission_list_api")
 ]

submission/views/oj.py

@@ -1,4 +1,3 @@
-from django_redis import get_redis_connection
 
 from account.decorators import login_required, check_contest_permission
 from judge.tasks import judge_task
@@ -9,6 +8,7 @@ from utils.throttling import TokenBucket, BucketController
 from ..models import Submission
 from ..serializers import CreateSubmissionSerializer, SubmissionModelSerializer
 from ..serializers import SubmissionSafeSerializer, SubmissionListSerializer
+from utils.cache import throttling_cache
 
 
 # from judge.dispatcher import JudgeDispatcher
@@ -17,7 +17,7 @@ from ..serializers import SubmissionSafeSerializer, SubmissionListSerializer
 def _submit(response, user, problem_id, language, code, contest_id):
     # TODO: 预设默认值，需修改
     controller = BucketController(user_id=user.id,
-                                  redis_conn=get_redis_connection("Throttling"),
+                                  redis_conn=throttling_cache,
                                   default_capacity=30)
     bucket = TokenBucket(fill_rate=10, capacity=20,
                          last_capacity=controller.last_capacity,

utils/cache.py

@@ -0,0 +1,6 @@
+from django.conf import settings
+from django_redis import get_redis_connection
+
+judge_queue_cache = get_redis_connection(settings.CACHE_JUDGE_QUEUE)
+throttling_cache = get_redis_connection(settings.CACHE_THROTTLING)
+default_cache = get_redis_connection("default")

utils/constants.py

@@ -0,0 +1,2 @@
+class CacheKey:
+    waiting_queue = "waiting_queue"