feat: 🚧 draft playground set

Author: Eason0729

backend/Cargo.lock

@@ -28,6 +28,8 @@ base64 = "0.21.5"
 uuid = "1.5.0"
 tonic-web = "0.9.2"
 quick_cache = "0.4.0"
+itertools = "0.12.0"
+leaky-bucket = "1.0.1"
 
 [dependencies.tokio-stream]
 version = "0.1.14"
@@ -43,7 +45,7 @@ features =["async-await", "log"]
 
 [dependencies.tokio]
 version = "1.28.0"
-features =["macros", "rt-multi-thread", "full"]
+features =["macros", "rt-multi-thread", "full","time"]
 
 [dependencies.sea-orm]
 version = "0.12.2"

backend/Cargo.toml

@@ -1,15 +1,20 @@
 use std::sync::Arc;
 
+use futures::Future;
+use leaky_bucket::RateLimiter;
 use sea_orm::{ActiveModelTrait, ActiveValue, EntityTrait, QueryOrder};
 use thiserror::Error;
 use tokio_stream::StreamExt;
 use tonic::Status;
 use uuid::Uuid;
 
 use crate::{
+    endpoint::util::{error, stream::TonicStream},
     grpc::{
-        backend::{submit_status, JudgeResult as BackendResult, SubmitStatus},
-        judger::{judge_response, JudgeRequest, JudgeResponse, JudgeResult, JudgerCode, TestIo},
+        backend::{submit_status, JudgeResult as BackendResult, PlaygroundResult, SubmitStatus},
+        judger::{
+            judge_response, JudgeRequest, JudgeResponse, JudgeResult, JudgerCode, LangInfo, TestIo,
+        },
     },
     init::{config::CONFIG, db::DB},
 };
@@ -21,7 +26,26 @@ use super::util::{
 };
 use entity::*;
 
-type TonicStream<T> = std::pin::Pin<Box<dyn tokio_stream::Stream<Item = Result<T, Status>> + Send>>;
+struct Waker;
+
+impl std::task::Wake for Waker {
+    fn wake(self: Arc<Self>) {
+        log::error!("waker wake");
+    }
+}
+
+macro_rules! check_rate_limit {
+    ($s:expr) => {{
+        let waker = Arc::new(Waker).into();
+        let mut cx = std::task::Context::from_waker(&waker);
+
+        let ac = $s.limiter.clone().acquire_owned(1);
+        tokio::pin!(ac);
+        if ac.as_mut().poll(&mut cx).is_pending() {
+            return Err(Error::RateLimit);
+        }
+    }};
+}
 
 #[derive(Debug, Error)]
 pub enum Error {
@@ -39,8 +63,8 @@ pub enum Error {
     Tonic(#[from] tonic::transport::Error),
     #[error("`{0}`")]
     Internal(&'static str),
-    // #[error("judger tls error")]
-    // TlsError,
+    #[error("Rate limit exceeded")]
+    RateLimit,
 }
 
 impl From<Error> for super::Error {
@@ -58,6 +82,9 @@ impl From<Error> for super::Error {
             Error::Tonic(x) => super::Error::Tonic(x),
             Error::Internal(x) => super::Error::Internal(x),
             Error::GrpcReport(x) => super::Error::GrpcReport(x),
+            Error::RateLimit => super::Error::GrpcReport(tonic::Status::resource_exhausted(
+                "judge rate limit exceeded",
+            )),
         }
     }
 }
@@ -93,10 +120,10 @@ impl From<JudgeResult> for SubmitStatus {
     }
 }
 
-#[derive(Clone)]
 pub struct SubmitController {
     router: Arc<Router>,
     pubsub: Arc<PubSub<Result<SubmitStatus, Status>, i32>>,
+    limiter: Arc<RateLimiter>,
 }
 
 impl SubmitController {
@@ -105,6 +132,14 @@ impl SubmitController {
         Ok(SubmitController {
             router: Router::new(&config.judger).await?,
             pubsub: Arc::new(PubSub::default()),
+            limiter: Arc::new(
+                RateLimiter::builder()
+                    .max(25)
+                    .initial(10)
+                    .refill(2)
+                    .interval(std::time::Duration::from_millis(100))
+                    .build(),
+            ),
         })
     }
     async fn stream(
@@ -170,6 +205,7 @@ impl SubmitController {
         }
     }
     pub async fn submit(&self, submit: Submit) -> Result<i32, Error> {
+        check_rate_limit!(self);
         let db = DB.get().unwrap();
 
         let mut conn = self.router.get(&submit.lang).await?;
@@ -225,6 +261,13 @@ impl SubmitController {
     pub async fn follow(&self, submit_id: i32) -> Option<TonicStream<SubmitStatus>> {
         self.pubsub.subscribe(&submit_id)
     }
+    pub fn list_lang(&self) -> Vec<LangInfo> {
+        self.router.langs()
+    }
+    pub async fn playground(&self) -> Result<TonicStream<PlaygroundResult>, Error> {
+        check_rate_limit!(self);
+        todo!()
+    }
 }
 
 impl From<Error> for Status {

backend/src/controller/submit.rs

@@ -7,6 +7,7 @@ use std::{
     },
 };
 
+use itertools::Itertools;
 use spin::{Mutex, RwLock};
 use tonic::*;
 use uuid::Uuid;
@@ -148,6 +149,9 @@ impl Upstream {
 
         self_
     }
+    fn langs(&self) -> Vec<LangInfo> {
+        self.langs.read().values().cloned().collect()
+    }
     async fn health_check(&self) -> Result<(), Error> {
         macro_rules! health {
             ($e:expr) => {
@@ -208,6 +212,14 @@ impl Router {
             next_entry: AtomicUsize::new(0),
         }))
     }
+    pub fn langs(&self) -> Vec<LangInfo> {
+        self.upstreams
+            .iter()
+            .map(|x| x.langs())
+            .flatten()
+            .unique_by(|x| x.lang_uid.clone())
+            .collect()
+    }
     pub async fn get(&self, uid: &Uuid) -> Result<ConnGuard, Error> {
         let server_count = self.upstreams.len();
         for _ in 0..(server_count * 2 + 1) {

backend/src/controller/util/router.rs

@@ -1,5 +1,6 @@
 pub mod contest;
 mod education;
+pub mod playground;
 pub mod problem;
 pub mod submit;
 pub mod testcase;

backend/src/endpoint/mod.rs

@@ -0,0 +1,19 @@
+use super::endpoints::*;
+use super::tools::*;
+use super::util::stream::TonicStream;
+
+use crate::grpc::backend::playground_set_server::*;
+use crate::grpc::backend::*;
+
+#[async_trait]
+impl PlaygroundSet for Arc<Server> {
+    #[doc = " Server streaming response type for the Run method."]
+    type RunStream = TonicStream<PlaygroundResult>;
+
+    async fn run(
+        &self,
+        req: Request<PlaygroundRequest>,
+    ) -> Result<Response<Self::RunStream>, Status> {
+        Err(Status::unimplemented("Not implemented"))
+    }
+}

backend/src/endpoint/playground.rs

@@ -8,6 +8,7 @@ use crate::controller::util::code::Code;
 use crate::grpc::backend::submit_set_server::*;
 use crate::grpc::backend::StateCode as BackendCode;
 use crate::grpc::backend::*;
+use crate::grpc::judger::LangInfo;
 
 use entity::{submit::*, *};
 use tokio_stream::wrappers::ReceiverStream;
@@ -249,4 +250,26 @@ impl SubmitSet for Arc<Server> {
 
         Ok(Response::new(()))
     }
+
+    #[doc = " Server streaming response type for the ListLangs method."]
+    type ListLangsStream = TonicStream<Language>;
+
+    async fn list_langs(&self, _: Request<()>) -> Result<Response<Self::ListLangsStream>, Status> {
+        let langs = self.submit.list_lang().into_iter().map(|x| Ok(x.into()));
+
+        Ok(Response::new(
+            Box::pin(tokio_stream::iter(langs)) as TonicStream<_>
+        ))
+    }
+}
+
+impl From<LangInfo> for Language {
+    fn from(value: LangInfo) -> Self {
+        Language {
+            lang_uid: value.lang_uid,
+            lang_name: value.lang_name,
+            info: value.info,
+            lang_ext: value.lang_ext,
+        }
+    }
 }

backend/src/endpoint/submit.rs

@@ -152,6 +152,13 @@ message RejudgeRequest {
   required string request_id = 2;
 }
 
+message Language {
+  required string lang_uid = 1;
+  required string lang_name = 2;
+  required string info = 3;
+  required string lang_ext = 4;
+}
+
 service SubmitSet {
   rpc List(ListRequest) returns (ListSubmitResponse);
   rpc Info(SubmitId) returns (SubmitInfo);
@@ -164,6 +171,8 @@ service SubmitSet {
   // are not guarantee to yield status
   rpc Follow(SubmitId) returns (stream SubmitStatus);
   rpc Rejudge(RejudgeRequest) returns (google.protobuf.Empty);
+
+  rpc ListLangs(google.protobuf.Empty) returns (stream Language);
 }
 
 // Announcements
@@ -615,3 +624,25 @@ service TokenSet {
 
   rpc Logout(google.protobuf.Empty) returns (google.protobuf.Empty);
 }
+
+message PlaygroundRequest{
+  required bytes code=1;
+  required string lang=2;
+  required bytes input=3;
+}
+
+message CompileLog{
+  required uint32 level = 1;
+  required string msg = 2;
+}
+
+message PlaygroundResult{
+  oneof result{
+    CompileLog compile=1;
+    bytes output=2;
+  }
+}
+
+service PlaygroundSet{
+  rpc Run(PlaygroundRequest) returns (stream PlaygroundResult);
+}