From 36c9f6611bd51a4163de033150f8ec61be1f5e4d Mon Sep 17 00:00:00 2001
From: masnn <masnn0@outlook.com>
Date: Thu, 22 Aug 2019 15:12:32 +0800
Subject: [PATCH] Fix #522

---
 vj4/handler/problem.py |  9 +++++++--
 vj4/model/domain.py    | 18 ++++++++++++++++++
 vj4/model/system.py    | 15 +++++++++++++++
 3 files changed, 40 insertions(+), 2 deletions(-)

diff --git a/vj4/handler/problem.py b/vj4/handler/problem.py
index e67030b8e..06d4bac62 100644
--- a/vj4/handler/problem.py
+++ b/vj4/handler/problem.py
@@ -584,8 +584,13 @@ async def post(self, *, title: str, content: str, hidden: bool=False, numeric_pi
     pid = None
     if numeric_pid:
       pid = await domain.inc_pid_counter(self.domain_id)
-    pid = await problem.add(self.domain_id, title, content, self.user['_id'],
-                            hidden=hidden, pid=pid)
+    try:
+      pid = await problem.add(self.domain_id, title, content, self.user['_id'],
+                              hidden=hidden, pid=pid)
+    except Exception as e:
+      if numeric_pid:
+        await domain.dec_pid_counter(self.domain_id)
+      raise
     self.json_or_redirect(self.reverse_url('problem_settings', pid=pid))
 
 
diff --git a/vj4/model/domain.py b/vj4/model/domain.py
index 4953d3bdc..26aa6cfbf 100644
--- a/vj4/model/domain.py
+++ b/vj4/model/domain.py
@@ -129,6 +129,24 @@ async def inc_pid_counter(domain_id: str):
   return doc['pid_counter']
 
 
+@argmethod.wrap
+async def dec_pid_counter(domain_id: str):
+  """Decrements the problem ID counter.
+
+  Returns:
+    Integer value before decrement.
+  """
+  for domain in builtin.DOMAINS:
+    if domain['_id'] == domain_id:
+      return await system.dec_pid_counter()
+  coll = db.coll('domain')
+  await coll.update_one(filter={'_id': domain_id, 'pid_counter': {'$exists': False}},
+                        update={'$set': {'pid_counter': 1000}})
+  doc = await coll.find_one_and_update(filter={'_id': domain_id},
+                                       update={'$dec': {'pid_counter': 1}})
+  return doc['pid_counter']
+
+
 @argmethod.wrap
 async def set_role(domain_id: str, role: str, perm: int):
   return await set_roles(domain_id, {role: perm})
diff --git a/vj4/model/system.py b/vj4/model/system.py
index 9a2bf65b6..70da57239 100644
--- a/vj4/model/system.py
+++ b/vj4/model/system.py
@@ -41,6 +41,21 @@ async def inc_pid_counter():
   return doc['value']
 
 
+@argmethod.wrap
+async def dec_pid_counter():
+  """Decrements the problem ID counter.
+
+  Returns:
+    Integer value before decrement.
+  """
+  coll = db.coll('system')
+  await coll.update_one(filter={'_id': 'pid_counter'},
+                        update={'$setOnInsert': {'value': 1000}}, upsert=True)
+  doc = await coll.find_one_and_update(filter={'_id': 'pid_counter'},
+                                       update={'$dec': {'value': 1}})
+  return doc['value']
+
+
 async def acquire_lock(lock_name: str):
   lock_value = random.randint(1, 0xFFFFFFFF)
   coll = db.coll('system')