use std::collections::{BTreeMap, BTreeSet};
use std::sync::Mutex;
use std::time::Duration;
use tokio::sync::{Notify, mpsc, oneshot};
use tracing::{info, warn};
use super::records::RecordSelection;
use super::{Engine, Occupant};
use crate::api;
use crate::cache;
use crate::config::Config;
use crate::docker::{DockerApi, Tty, TtySize};
use crate::github::GithubApi;
use crate::isolation;
use crate::scheduler::SlotRef;
use crate::store::Store;
use crate::sync::lock;
pub(super) const CANCEL_GRACE: Duration = Duration::from_secs(90);
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum Shutdown {
No,
Graceful,
Force,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum RunningJobs {
Keep,
Kill,
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum Target {
Repository(String),
All,
}
#[derive(Debug, Clone, PartialEq, Eq, thiserror::Error)]
pub enum ControlError {
#[error("{0}")]
NotFound(String),
#[error("{message}")]
Busy {
message: String,
jobs: Vec<api::RunningJob>,
},
#[error("{0}")]
Conflict(String),
#[error("{0}")]
Invalid(String),
#[error("{0}")]
Failed(String),
#[error("{0}")]
NotRecorded(String),
#[error("runnerd is stopping")]
Stopped,
}
#[derive(Debug)]
pub enum Command {
Pause {
target: Target,
jobs: RunningJobs,
reply: oneshot::Sender<Result<api::Paused, ControlError>>,
},
Resume {
target: Target,
reply: oneshot::Sender<Result<(), ControlError>>,
},
Reseed {
profile: String,
reply: oneshot::Sender<Result<(), ControlError>>,
},
CleanCache {
name: String,
reply: oneshot::Sender<Result<(), ControlError>>,
},
ReplaceConfig {
config: Box<Config>,
memory_max: Option<u64>,
commit: Commit,
reply: oneshot::Sender<Result<Option<String>, ControlError>>,
},
RemoveRecords {
selection: RecordSelection,
reply: oneshot::Sender<Result<api::RemovedRecords, ControlError>>,
},
OpenShell {
job: u64,
size: TtySize,
reply: oneshot::Sender<Result<Tty, ControlError>>,
},
CloseShell {
container: String,
},
}
pub struct Commit(Box<dyn FnOnce() -> Result<Option<String>, String> + Send>);
impl Commit {
pub fn new(write: impl FnOnce() -> Result<Option<String>, String> + Send + 'static) -> Commit {
Commit(Box::new(write))
}
pub async fn run(self) -> Result<Option<String>, String> {
tokio::task::spawn_blocking(self.0)
.await
.unwrap_or_else(|join_error| Err(join_error.to_string()))
}
}
impl std::fmt::Debug for Commit {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
f.write_str("Commit")
}
}
#[derive(Debug)]
pub struct Commands(mpsc::UnboundedReceiver<Command>);
impl Commands {
pub async fn recv(&mut self) -> Option<Command> {
self.0.recv().await
}
}
#[derive(Debug)]
pub struct Control {
repositories: Mutex<BTreeSet<String>>,
paused: Mutex<BTreeSet<String>>,
store: Store,
shutdown: Mutex<Shutdown>,
pub(super) wake: Notify,
commands: mpsc::UnboundedSender<Command>,
}
impl Control {
pub fn new(
repositories: impl IntoIterator<Item = String>,
paused: impl IntoIterator<Item = String>,
store: Store,
) -> (Control, Commands) {
let repositories: BTreeSet<String> = repositories.into_iter().collect();
let paused: BTreeSet<String> = paused
.into_iter()
.filter(|repository| repositories.contains(repository))
.collect();
drop(store.paused(paused.iter().cloned().collect()));
let (commands, receiver) = mpsc::unbounded_channel();
let control = Control {
repositories: Mutex::new(repositories),
paused: Mutex::new(paused),
store,
shutdown: Mutex::new(Shutdown::No),
wake: Notify::new(),
commands,
};
(control, Commands(receiver))
}
pub fn mode(&self) -> Shutdown {
*lock(&self.shutdown)
}
pub async fn set_paused(&self, target: &Target, paused: bool) -> Result<(), ControlError> {
let recorded = {
let repositories = lock(&self.repositories);
let names: Vec<&String> = match target {
Target::Repository(repository) => {
let Some(name) = repositories.get(repository) else {
return Err(not_configured(repository, repositories.iter()));
};
vec![name]
}
Target::All => repositories.iter().collect(),
};
let mut set = lock(&self.paused);
for name in names {
if paused {
set.insert(name.clone());
} else {
set.remove(name);
}
}
self.store.paused(set.iter().cloned().collect())
};
self.wake.notify_one();
if recorded.await {
return Ok(());
}
let name = match target {
Target::Repository(repository) => repository.as_str(),
Target::All => "every repository",
};
let (done, undone) = if paused {
("paused", "resumes it")
} else {
("resumed", "pauses it again")
};
let message = format!(
"{name} is {done} until runnerd restarts, but the job history cannot record it, so a restart {undone}"
);
warn!("{message}");
Err(ControlError::NotRecorded(message))
}
fn set_repositories(&self, repositories: BTreeSet<String>) {
let mut configured = lock(&self.repositories);
let mut paused = lock(&self.paused);
let before = paused.len();
paused.retain(|repository| repositories.contains(repository));
if paused.len() != before {
drop(self.store.paused(paused.iter().cloned().collect()));
}
*configured = repositories;
}
pub fn shutdown(&self, jobs: RunningJobs) {
let mode = match jobs {
RunningJobs::Keep => Shutdown::Graceful,
RunningJobs::Kill => Shutdown::Force,
};
let mut shutdown = lock(&self.shutdown);
*shutdown = (*shutdown).max(mode);
drop(shutdown);
self.wake.notify_one();
}
pub fn is_paused(&self, repository: &str) -> bool {
lock(&self.paused).contains(repository)
}
async fn send<T>(
&self,
command: impl FnOnce(oneshot::Sender<T>) -> Command,
) -> Result<T, ControlError> {
let (reply, answer) = oneshot::channel();
self.commands
.send(command(reply))
.map_err(|_| ControlError::Stopped)?;
answer.await.map_err(|_| ControlError::Stopped)
}
pub async fn pause(
&self,
target: Target,
jobs: RunningJobs,
) -> Result<api::Paused, ControlError> {
self.send(|reply| Command::Pause {
target,
jobs,
reply,
})
.await?
}
pub async fn resume(&self, target: Target) -> Result<(), ControlError> {
self.send(|reply| Command::Resume { target, reply }).await?
}
pub async fn reseed(&self, profile: &str) -> Result<(), ControlError> {
self.send(|reply| Command::Reseed {
profile: profile.to_owned(),
reply,
})
.await?
}
pub async fn clean_cache(&self, name: &str) -> Result<(), ControlError> {
self.send(|reply| Command::CleanCache {
name: name.to_owned(),
reply,
})
.await?
}
pub async fn remove_records(
&self,
selection: RecordSelection,
) -> Result<api::RemovedRecords, ControlError> {
self.send(|reply| Command::RemoveRecords { selection, reply })
.await?
}
pub async fn open_shell(&self, job: u64, size: TtySize) -> Result<Tty, ControlError> {
self.send(|reply| Command::OpenShell { job, size, reply })
.await?
}
pub fn close_shell(&self, container: String) {
let _ = self.commands.send(Command::CloseShell { container });
}
pub async fn replace_config(
&self,
config: Config,
memory_max: Option<u64>,
commit: Commit,
) -> Result<Option<String>, ControlError> {
self.send(|reply| Command::ReplaceConfig {
config: Box::new(config),
memory_max,
commit,
reply,
})
.await?
}
}
fn not_configured<'name>(
repository: &str,
configured: impl Iterator<Item = &'name String>,
) -> ControlError {
let configured: Vec<&str> = configured.map(String::as_str).collect();
let configured = if configured.is_empty() {
"none is configured".to_owned()
} else {
format!("configured: {}", configured.join(", "))
};
ControlError::NotFound(format!(
"{repository} is not a configured repository; {configured}"
))
}
fn list(jobs: &[api::RunningJob]) -> String {
jobs.iter()
.map(ToString::to_string)
.collect::<Vec<_>>()
.join(", ")
}
impl<G: GithubApi + 'static, D: DockerApi + 'static> Engine<G, D> {
pub(super) async fn command(&mut self, command: Command) {
match command {
Command::Pause {
target,
jobs,
reply,
} => {
let _ = reply.send(self.pause(&target, jobs).await);
}
Command::Resume { target, reply } => {
let _ = reply.send(self.control.set_paused(&target, false).await);
}
Command::Reseed { profile, reply } => {
let _ = reply.send(self.request_reseed(&profile));
}
Command::CleanCache { name, reply } => {
let _ = reply.send(self.clean_cache(&name).await);
}
Command::ReplaceConfig {
config,
memory_max,
commit,
reply,
} => {
let answer = match commit.run().await {
Ok(warning) => {
self.reload(*config, memory_max);
Ok(warning)
}
Err(message) => Err(ControlError::Failed(message)),
};
let _ = reply.send(answer);
}
Command::RemoveRecords { selection, reply } => {
let _ = reply.send(self.remove_records(&selection).await);
}
Command::OpenShell { job, size, reply } => {
let opened = self.open_shell(job, size).await;
if let Err(Ok(tty)) = reply.send(opened) {
self.close_shell(&tty.container).await;
}
}
Command::CloseShell { container } => self.close_shell(&container).await,
}
}
pub(super) async fn pause(
&mut self,
target: &Target,
running: RunningJobs,
) -> Result<api::Paused, ControlError> {
let repositories: BTreeSet<String> = match target {
Target::Repository(repository) if !self.repositories.contains_key(repository) => {
return Err(not_configured(repository, self.repositories.keys()));
}
Target::Repository(repository) => BTreeSet::from([repository.clone()]),
Target::All => self.repositories.keys().cloned().collect(),
};
self.take_found();
let runners: Vec<(SlotRef, Occupant)> = self
.occupied()
.filter(|(_, occupant)| repositories.contains(&occupant.repository))
.map(|(slot, occupant)| (slot, occupant.clone()))
.collect();
let jobs: Vec<api::RunningJob> = runners
.iter()
.filter(|(_, occupant)| occupant.busy || occupant.job.is_some())
.map(|(_, occupant)| occupant.running_job())
.collect();
if running == RunningJobs::Keep && !jobs.is_empty() {
let mut by_repository: BTreeMap<&str, Vec<api::RunningJob>> = BTreeMap::new();
for job in &jobs {
by_repository
.entry(&job.repository)
.or_default()
.push(job.clone());
}
let runs: Vec<String> = by_repository
.iter()
.map(|(repository, jobs)| format!("{repository} runs {}", list(jobs)))
.collect();
return Err(ControlError::Busy {
message: runs.join("; "),
jobs,
});
}
let recorded = self.control.set_paused(target, true).await;
match target {
Target::Repository(repository) => info!("{repository}: paused"),
Target::All => info!("paused every repository"),
}
let mut stopped = Vec::new();
if running == RunningJobs::Keep {
return recorded.map(|()| api::Paused { stopped });
}
for (slot, occupant) in runners {
let repository = &occupant.repository;
let cancelled = if let Some(job) = occupant.job {
self.cancel(repository, job).await
} else {
if occupant.busy {
warn!(
"{repository}: {} runs a job its log has not named yet; GitHub fails it once the runner is gone",
occupant.runner_name
);
}
false
};
if cancelled {
self.cancelling.insert(
occupant.runner_name.clone(),
tokio::time::Instant::now() + CANCEL_GRACE,
);
} else {
self.stop(&slot, RunningJobs::Kill).await;
if self
.occupied()
.any(|(_, candidate)| candidate.runner_name == occupant.runner_name)
{
continue;
}
}
stopped.push(occupant.running_job());
}
match recorded {
Ok(()) => Ok(api::Paused { stopped }),
Err(ControlError::NotRecorded(message)) if !stopped.is_empty() => Err(
ControlError::NotRecorded(format!("{message}; stopped {}", list(&stopped))),
),
Err(error) => Err(error),
}
}
pub(super) async fn kill_cancelled(&mut self) {
let now = tokio::time::Instant::now();
let running: BTreeSet<String> = self
.occupied()
.map(|(_, occupant)| occupant.runner_name.clone())
.collect();
self.cancelling.retain(|name, _| running.contains(name));
let overdue: Vec<(SlotRef, String)> = self
.occupied()
.filter(|(_, occupant)| {
self.cancelling
.get(&occupant.runner_name)
.is_some_and(|deadline| *deadline <= now)
})
.map(|(slot, occupant)| (slot, occupant.runner_name.clone()))
.collect();
for (slot, name) in overdue {
warn!(
"{name} did not exit {}s after its run was cancelled",
CANCEL_GRACE.as_secs()
);
self.stop(&slot, RunningJobs::Kill).await;
}
}
async fn cancel(&self, repository: &str, job: u64) -> bool {
let run = match self.github.job(repository, job).await {
Ok(details) => details.run_id,
Err(error) => {
warn!("{repository}: cannot read job {job} to cancel its run: {error}");
return false;
}
};
match self.github.cancel_run(repository, run).await {
Ok(()) => {
info!("{repository}: cancelled run {run} of job {job}");
true
}
Err(error) => {
warn!("{repository}: cannot cancel run {run}: {error}");
false
}
}
}
pub(super) fn request_reseed(&mut self, profile: &str) -> Result<(), ControlError> {
if !self.config.profiles.contains_key(profile) {
return Err(ControlError::NotFound(format!(
"{profile} is not a configured profile"
)));
}
if self.is_reseeding(profile) {
return Err(ControlError::Conflict(format!(
"{profile} is reseeding already"
)));
}
if self.shell_profiles().contains(profile) {
return Err(ControlError::Conflict(format!(
"a shell is open in a record of {profile}, which mounts its snapshot; reseed it once the shell exits"
)));
}
info!("{profile}: reseeding once its slots are free");
self.reseed_requested.insert(profile.to_owned());
Ok(())
}
pub(super) async fn clean_cache(&mut self, name: &str) -> Result<(), ControlError> {
let Some(volume) = cache::volumes(&self.names, &self.config)
.into_iter()
.find(|volume| volume.name == name)
else {
return Err(ControlError::NotFound(format!(
"{name} is not a cache volume of this config; `runnerctl caches` lists them"
)));
};
let users: Vec<api::RunningJob> = match &volume.kind {
api::CacheKind::Snapshot { profile } => {
return Err(ControlError::Invalid(format!(
"{name} is {profile}'s snapshot, which every slot's layers sit on; reseed {profile} to rebuild it"
)));
}
api::CacheKind::Slot { profile, slot } => {
if self.reseeds.contains_key(profile) {
return Err(ControlError::Conflict(format!(
"{profile} is reseeding, which writes {name}"
)));
}
self.occupied()
.filter(|(at, _)| at.profile == *profile && at.index + 1 == *slot)
.map(|(_, occupant)| occupant.running_job())
.collect()
}
api::CacheKind::Pnpm | api::CacheKind::Maven | api::CacheKind::Mirror => self
.occupied()
.map(|(_, occupant)| occupant.running_job())
.collect(),
};
let shared = matches!(volume.kind, api::CacheKind::Pnpm | api::CacheKind::Maven);
if shared && !self.records.shells.is_empty() {
return Err(ControlError::Conflict(format!(
"a shell open in a record mounts {name}; clean it once the shell exits"
)));
}
if !users.is_empty() {
return Err(ControlError::Busy {
message: format!(
"{name} is in use by {}; clean it once they finish",
list(&users)
),
jobs: users,
});
}
if volume.kind == api::CacheKind::Mirror && self.mirror_up {
cache::mirror_down(&*self.docker, &self.names)
.await
.map_err(|error| {
ControlError::Failed(format!("cannot stop the mirror: {error}"))
})?;
self.mirror_up = false;
}
self.docker
.remove_volume(name)
.await
.map_err(|error| ControlError::Failed(format!("cannot remove {name}: {error}")))?;
info!("removed {name}");
self.trim_due = true;
Ok(())
}
pub(super) fn reload(&mut self, config: Config, memory_max: Option<u64>) {
let mut old = std::mem::replace(&mut self.config, config);
for (name, profile) in &self.config.profiles {
let slots = self.slots.entry(name.clone()).or_default();
if slots.len() < profile.slots as usize {
slots.resize(profile.slots as usize, None);
}
self.retired.remove(name);
}
for (name, profile) in std::mem::take(&mut old.profiles) {
if !self.config.profiles.contains_key(&name) {
self.retired.insert(name, profile);
}
}
self.drop_closed_slots();
let names: BTreeSet<String> = self
.config
.repositories
.iter()
.map(|repository| repository.name.clone())
.collect();
let removed: Vec<String> = self
.repositories
.keys()
.filter(|name| !names.contains(*name))
.cloned()
.collect();
self.draining.extend(removed);
self.draining.retain(|name| !names.contains(name));
self.repositories.retain(|name, _| names.contains(name));
for name in &names {
self.repositories.entry(name.clone()).or_default();
}
self.control.set_repositories(names);
let profiles = &self.config.profiles;
self.reseed_requested
.retain(|name| profiles.contains_key(name));
self.isolation = isolation::backend(self.config.host.isolation);
self.memory_max = memory_max;
self.caches
.send_replace(cache::volumes(&self.names, &self.config));
info!("reloaded the config");
}
}