Issue23 (#28)

* feat(Testsuit): ✨ make testsuit an cli application instead of a collection of rstest * fix(Testsuit): 🐛 bring nighty frontend * feat(Testsuit): ✅ add create and add_to test * test(Testsuit): 🧪 lua RE(unknown) * fix(Judger): 🐛 change return code form 1 to 0 * feat(Backend): ✨ auto generate config and database * cargo clippy * refactor(Backend): 🚚 move OnceCell<DB> to field in struct * cargo fmt * ci(CI): 🚀 update staging.yml to build docker container * merge back to master (#27) * ci(CI): 🚀 fix CI * feat(Testsuit): ✨ make testsuit an cli application instead of a collection of rstest * fix(Testsuit): 🐛 bring nighty frontend * feat(Testsuit): ✅ add create and add_to test * test(Testsuit): 🧪 lua RE(unknown) * fix(Judger): 🐛 change return code form 1 to 0 * feat(Backend): ✨ auto generate config and database * cargo clippy * regen lockfile
mdcpp · Jan 28, 2024 · 1f5bd8d · 1f5bd8d
1 parent f4a6730
commit 1f5bd8d
Show file tree

Hide file tree

Showing 30 changed files with 682 additions and 532 deletions.
diff --git a/Cargo.lock b/Cargo.lock
diff --git a/backend/justfile b/backend/justfile
@@ -4,13 +4,8 @@
 #     DATABASE_URL=sqlite://backend.sqlite sea-orm-cli migrate refresh
 #     sea-orm-cli generate entity -u sqlite://backend.sqlite -o entities
 
-regen-database:
-    rm database/*
-    sea-orm-cli migrate -u sqlite://database/backend.sqlite?mode=rwc
-
 prepare:
     mkdir -p database
-    mkdir -p config
     sea-orm-cli migrate -u sqlite://database/backend.sqlite?mode=rwc
 
 entity-codegen:

diff --git a/backend/src/controller/judger/mod.rs b/backend/src/controller/judger/mod.rs
@@ -2,9 +2,12 @@ mod pubsub;
 mod route;
 mod score;
 
-use std::sync::{
-    atomic::{AtomicI64, Ordering},
-    Arc,
+use std::{
+    ops::Deref,
+    sync::{
+        atomic::{AtomicI64, Ordering},
+        Arc,
+    },
 };
 use tokio_stream::StreamExt;
 
@@ -16,18 +19,15 @@ use crate::{
 use futures::Future;
 use leaky_bucket::RateLimiter;
 use opentelemetry::{global, metrics::ObservableGauge};
-use sea_orm::{ActiveModelTrait, ActiveValue, EntityTrait, QueryOrder};
+use sea_orm::{ActiveModelTrait, ActiveValue, DatabaseConnection, EntityTrait, QueryOrder};
 use thiserror::Error;
 use tonic::Status;
 use tracing::{instrument, Instrument, Span};
 use uuid::Uuid;
 
-use crate::{
-    grpc::{
-        backend::{submit_status, JudgeResult as BackendResult, PlaygroundResult, SubmitStatus},
-        judger::*,
-    },
-    init::db::DB,
+use crate::grpc::{
+    backend::{submit_status, JudgeResult as BackendResult, PlaygroundResult, SubmitStatus},
+    judger::*,
 };
 
 use self::{
@@ -157,11 +157,16 @@ pub struct JudgerController {
     pubsub: Arc<PubSub<Result<SubmitStatus, Status>, i32>>,
     limiter: Arc<RateLimiter>,
     running_meter: (AtomicI64, ObservableGauge<i64>),
+    db: Arc<DatabaseConnection>,
 }
 
 impl JudgerController {
     #[tracing::instrument(parent=span, name="judger_construct",level = "info",skip_all)]
-    pub async fn new(config: Vec<config::Judger>, span: &Span) -> Result<Self, Error> {
+    pub async fn new(
+        config: Vec<config::Judger>,
+        db: Arc<DatabaseConnection>,
+        span: &Span,
+    ) -> Result<Self, Error> {
         let router = Router::new(config, span)?;
         Ok(JudgerController {
             router,
@@ -180,6 +185,7 @@ impl JudgerController {
                     .i64_observable_gauge("running_judge")
                     .init(),
             ),
+            db,
         })
     }
     fn record(&self) -> MeterGuard {
@@ -257,14 +263,14 @@ impl JudgerController {
         model.memory = ActiveValue::Set(Some(mem.try_into().unwrap_or(i64::MAX)));
 
         model
-            .update(DB.get().unwrap())
+            .update(self.db.deref())
             .in_current_span()
             .await
             .map_err(Into::<Error>::into)
     }
     pub async fn submit(self: &Arc<Self>, req: Submit) -> Result<i32, Error> {
         check_rate_limit!(self);
-        let db = DB.get().unwrap();
+        let db = self.db.deref();
 
         let mut binding = problem::Entity::find_by_id(req.problem)
             .find_with_related(test::Entity)
@@ -315,14 +321,15 @@ impl JudgerController {
         conn.report_success();
 
         let self_ = self.clone();
+        let db = self.db.clone();
         tokio::spawn(async move {
             match self_
                 .stream(tx, res.into_inner(), submit_model, scores, submit_id)
                 .await
             {
                 Ok(submit) => {
                     score::ScoreUpload::new(req.user, problem, submit.score)
-                        .upload()
+                        .upload(&db)
                         .await;
                 }
                 Err(err) => {

diff --git a/backend/src/controller/judger/score.rs b/backend/src/controller/judger/score.rs
@@ -1,11 +1,11 @@
 use super::{user_contest, DebugName};
 use crate::{
     entity::{contest, problem, user},
-    init::db::DB,
     util::error::Error,
 };
 use sea_orm::{
-    ActiveModelTrait, ActiveValue, EntityTrait, IntoActiveModel, ModelTrait, TransactionTrait,
+    ActiveModelTrait, ActiveValue, DatabaseConnection, EntityTrait, IntoActiveModel, ModelTrait,
+    TransactionTrait,
 };
 use tracing::instrument;
 
@@ -27,10 +27,10 @@ impl ScoreUpload {
         }
     }
     #[instrument(skip(self))]
-    pub async fn upload(self) {
+    pub async fn upload(self, db: &DatabaseConnection) {
         let self_ = self;
         let mut retries = MAX_RETRY;
-        while let Err(err) = self_.upload_contest().await {
+        while let Err(err) = self_.upload_contest(db).await {
             tracing::debug!(err = err.to_string(), "retry_upload");
             match retries.checked_sub(1) {
                 None => {
@@ -42,7 +42,7 @@ impl ScoreUpload {
                 }
             }
         }
-        while let Err(err) = self_.upload_user().await {
+        while let Err(err) = self_.upload_user(db).await {
             tracing::debug!(err = err.to_string(), "retry_upload");
             match retries.checked_sub(1) {
                 None => {
@@ -55,9 +55,7 @@ impl ScoreUpload {
             }
         }
     }
-    async fn upload_user(&self) -> Result<(), Error> {
-        let db = DB.get().unwrap();
-
+    async fn upload_user(&self, db: &DatabaseConnection) -> Result<(), Error> {
         let txn = db.begin().await?;
 
         if self.user_id == self.problem.user_id {
@@ -83,9 +81,7 @@ impl ScoreUpload {
 
         txn.commit().await.map_err(Into::<Error>::into)
     }
-    async fn upload_contest(&self) -> Result<(), Error> {
-        let db = DB.get().unwrap();
-
+    async fn upload_contest(&self, db: &DatabaseConnection) -> Result<(), Error> {
         let txn = db.begin().await?;
 
         if self.user_id == self.problem.user_id {

diff --git a/backend/src/controller/token.rs b/backend/src/controller/token.rs
@@ -3,13 +3,15 @@ use chrono::{Duration, Local, NaiveDateTime};
 use quick_cache::sync::Cache;
 use rand::{Rng, SeedableRng};
 use rand_hc::Hc128Rng;
-use sea_orm::{ActiveModelTrait, ActiveValue, ColumnTrait, EntityTrait, QueryFilter};
+use sea_orm::{
+    ActiveModelTrait, ActiveValue, ColumnTrait, DatabaseConnection, EntityTrait, QueryFilter,
+};
 use spin::Mutex;
-use std::sync::Arc;
+use std::{ops::Deref, sync::Arc};
 use tokio::time;
 use tracing::{instrument, Instrument, Span};
 
-use crate::{init::db::DB, report_internal};
+use crate::report_internal;
 
 use super::metrics::RateMetrics;
 
@@ -64,27 +66,28 @@ pub struct TokenController {
     cache: Cache<Rand, CachedToken>,
     rng: Mutex<Hc128Rng>,
     cache_meter: RateMetrics<30>,
+    db: Arc<DatabaseConnection>,
 }
 
 impl TokenController {
     #[tracing::instrument(parent = span,name="token_construct_controller",level = "info",skip_all)]
-    pub fn new(span: &Span) -> Arc<Self> {
+    pub fn new(span: &Span, db: Arc<DatabaseConnection>) -> Arc<Self> {
         log::debug!("Setup TokenController");
         let cache = Cache::new(CACHE_SIZE);
         let self_ = Arc::new(Self {
             cache,
             rng: Mutex::new(Hc128Rng::from_entropy()),
             cache_meter: RateMetrics::new("hitrate_token"),
+            db: db.clone(),
         });
         tokio::spawn(async move {
-            let db = DB.get().unwrap();
             loop {
                 time::sleep(CLEAN_DUR).await;
                 let now = Local::now().naive_local();
 
                 if let Err(err) = token::Entity::delete_many()
                     .filter(token::Column::Expiry.lte(now))
-                    .exec(db)
+                    .exec(db.deref())
                     .await
                 {
                     log::error!("Token clean failed: {}", err);
@@ -99,8 +102,6 @@ impl TokenController {
         user: &crate::entity::user::Model,
         dur: Duration,
     ) -> Result<(String, NaiveDateTime), Error> {
-        let db = DB.get().unwrap();
-
         let rand: Rand = { self.rng.lock().gen() };
 
         let expiry = (Local::now() + dur).naive_local();
@@ -112,7 +113,7 @@ impl TokenController {
             permission: ActiveValue::Set(user.permission),
             ..Default::default()
         }
-        .insert(db)
+        .insert(self.db.deref())
         .in_current_span()
         .await?;
 
@@ -125,7 +126,6 @@ impl TokenController {
     #[instrument(skip_all, name = "token_verify_controller", level = "debug")]
     pub async fn verify(&self, token: &str) -> Result<(i32, RoleLv), Error> {
         let now = Local::now().naive_local();
-        let db = DB.get().unwrap();
 
         let rand =
             base64::Engine::decode(&base64::engine::general_purpose::STANDARD_NO_PAD, token)?;
@@ -156,7 +156,7 @@ impl TokenController {
             None => {
                 token = (token::Entity::find()
                     .filter(token::Column::Rand.eq(rand.to_vec()))
-                    .one(db)
+                    .one(self.db.deref())
                     .in_current_span()
                     .await?
                     .ok_or(Error::NonExist)?)
@@ -180,15 +180,13 @@ impl TokenController {
     }
     #[instrument(skip_all, name="token_remove_controller",level="debug", fields(token = token))]
     pub async fn remove(&self, token: String) -> Result<Option<()>, Error> {
-        let db = DB.get().unwrap();
-
         let rand =
             base64::Engine::decode(&base64::engine::general_purpose::STANDARD_NO_PAD, token)?;
         let rand: Rand = rand.try_into().map_err(|_| Error::InvalidTokenLength)?;
 
         token::Entity::delete_many()
             .filter(token::Column::Rand.eq(rand.to_vec()))
-            .exec(db)
+            .exec(self.db.deref())
             .await?;
 
         self.cache.remove(&rand);
@@ -204,7 +202,7 @@ impl TokenController {
         user_id: i32,
         // txn: &DatabaseTransaction,
     ) -> Result<(), Error> {
-        let db = DB.get().unwrap();
+        let db = self.db.deref();
 
         let models = token::Entity::find()
             .filter(token::Column::UserId.eq(user_id))