Something went wrong. Try again.
A Forgejo Runner plugin that spawn jobs as Kubernetes pods
Something went wrong. Try again.
123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257use std::{collections::HashMap, env};
use kube::Client;use tokio::sync::{mpsc, oneshot};use tokio_stream::wrappers::ReceiverStream;use tonic::{Request, Response, Status};use tracing::warn;
use crate::{ config::Config, plugin::{ CapabilitiesRequest, CapabilitiesResponse, CopyInChunk, CopyInResponse, CopyOutChunk, CopyOutRequest, CreateRequest, CreateResponse, ExecOutput, ExecRequest, RemoveRequest, RemoveResponse, StartComplete, StartOutput, StartRequest, backend_plugin_server::BackendPlugin, start_output, },};
mod environment;mod pod;mod validation;mod workflow;
use self::{ environment::JobEnvironments, pod::{ROOT_PATH, TEMP_PATH}, workflow::WorkflowOperations,};
const ACT_PATH: &str = "/shared/act";const TOOL_CACHE_PATH: &str = "/shared/toolcache";const DEFAULT_WORKDIR: &str = "/shared/workdir";const DEFAULT_PATH: &str = "/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin";
#[derive(Clone)]pub struct PluginServer { environments: JobEnvironments, workflow: WorkflowOperations,}
impl PluginServer { pub fn new(client: Client, config: Config) -> Self { let environments = JobEnvironments::new(client, config); Self { workflow: WorkflowOperations::new(environments.clone()), environments, } }
fn response_for(id: String) -> CreateResponse { CreateResponse { environment_id: id, root_path: ROOT_PATH.to_owned(), act_path: ACT_PATH.to_owned(), tool_cache_path: TOOL_CACHE_PATH.to_owned(), temp_path: TEMP_PATH.to_owned(), path_variable_name: Some("PATH".to_owned()), default_path_variable: Some(DEFAULT_PATH.to_owned()), path_separator: Some(":".to_owned()), environment_case_insensitive: false, os: "Linux".to_owned(), arch: env::consts::ARCH.to_owned(), } }
fn exec_cancellation_guard<T>( &self, id: String, sender: mpsc::Sender<Result<T, Status>>, ) -> oneshot::Sender<()> where T: Send + 'static, { let (finished, mut finished_rx) = oneshot::channel(); let server = self.clone(); tokio::spawn(async move { let cancelled = tokio::select! { // Completion precedes terminal publication, so it wins even if // the runner has already received that message and closed. biased; result = &mut finished_rx => result.is_err(), _ = sender.closed() => true, }; if cancelled { warn!(environment_id = %id, "Exec ended without completion; deleting job pod"); if let Err(error) = server.environments.cancel(&id).await { warn!(environment_id = %id, error = %error, "could not delete cancelled job pod"); } } }); finished }
async fn run_exec( self, request: ExecRequest, sender: mpsc::Sender<Result<ExecOutput, Status>>, ) { let id = request.environment_id.clone(); let mut completion = Some(self.exec_cancellation_guard(id.clone(), sender.clone())); let result = tokio::select! { _ = sender.closed() => Err(Status::cancelled("Exec stream closed")), result = self.workflow.exec(request, sender.clone(), &mut completion) => result, }; if let Some(completion) = completion && !sender.is_closed() { let _ = completion.send(()); } if let Err(error) = result { let _ = sender.send(Err(error)).await; } }}
#[tonic::async_trait]impl BackendPlugin for PluginServer { type StartStream = ReceiverStream<Result<StartOutput, Status>>; type ExecStream = ReceiverStream<Result<ExecOutput, Status>>; type CopyOutStream = ReceiverStream<Result<CopyOutChunk, Status>>;
async fn capabilities( &self, _request: Request<CapabilitiesRequest>, ) -> Result<Response<CapabilitiesResponse>, Status> { Ok(Response::new(CapabilitiesResponse { name: "kubernetes".to_owned(), })) }
async fn create( &self, request: Request<CreateRequest>, ) -> Result<Response<CreateResponse>, Status> { let request = request.into_inner(); let id = self.environments.create(&request).await?; Ok(Response::new(Self::response_for(id))) }
async fn start( &self, request: Request<StartRequest>, ) -> Result<Response<Self::StartStream>, Status> { let id = request.into_inner().environment_id; self.environments.ensure_exists(&id).await?; let (sender, receiver) = mpsc::channel(16); let server = self.clone(); tokio::spawn(async move { let _guard = match server.environments.acquire(&id).await { Ok(guard) => guard, Err(error) => { let _ = sender.send(Err(error)).await; return; } }; let ready = tokio::select! { biased; _ = sender.closed() => { if let Err(error) = server.environments.cancel(&id).await { let _ = sender.send(Err(error)).await; } return; } result = server.environments.wait_until_started(&id) => result, }; match ready { Ok(()) => { let output = StartOutput { output: Some(start_output::Output::StartComplete(StartComplete { image_env: HashMap::from([( "PATH".to_owned(), DEFAULT_PATH.to_owned(), )]), })), }; if sender.send(Ok(output)).await.is_err() { let _ = sender .send(Err(Status::cancelled("Start stream closed"))) .await; } } Err(error) => { warn!(environment_id = %id, pod = %environment::JobEnvironments::pod_name(&id), error = %error, "job pod failed to become ready"); let _ = sender.send(Err(error)).await; } } }); Ok(Response::new(ReceiverStream::new(receiver))) }
async fn exec( &self, request: Request<ExecRequest>, ) -> Result<Response<Self::ExecStream>, Status> { let request = request.into_inner(); environment::validate_environment_id(&request.environment_id)?; workflow::validate_exec_request(&request)?; self.environments .ensure_exists(&request.environment_id) .await?; let (sender, receiver) = mpsc::channel(32); tokio::spawn(self.clone().run_exec(request, sender)); Ok(Response::new(ReceiverStream::new(receiver))) }
async fn copy_in( &self, request: Request<tonic::Streaming<CopyInChunk>>, ) -> Result<Response<CopyInResponse>, Status> { self.workflow.copy_in(request.into_inner()).await?; Ok(Response::new(CopyInResponse {})) }
async fn copy_out( &self, request: Request<CopyOutRequest>, ) -> Result<Response<Self::CopyOutStream>, Status> { let request = request.into_inner(); self.environments .ensure_exists(&request.environment_id) .await?; let (sender, receiver) = mpsc::channel(16); tokio::spawn(self.workflow.clone().run_copy_out(request, sender)); Ok(Response::new(ReceiverStream::new(receiver))) }
async fn remove( &self, request: Request<RemoveRequest>, ) -> Result<Response<RemoveResponse>, Status> { let id = request.into_inner().environment_id; self.environments.remove(&id).await?; Ok(Response::new(RemoveResponse {})) }}
#[cfg(test)]#[path = "server_tests.rs"]mod supervision_tests;
#[cfg(test)]mod tests { use super::*;
#[test] fn exec_environment_supports_action_input_keys() { let request = ExecRequest { environment_id: "kube-396ba2f236d4d2462e9368b5".to_owned(), command: vec!["node".to_owned(), "action.js".to_owned()], env: HashMap::from([("INPUT_SET-SAFE-DIRECTORY".to_owned(), "true".to_owned())]), user: None, workdir: "/shared/workdir".to_owned(), }; workflow::validate_exec_request(&request).unwrap(); }}