mod serve;
use std::collections::HashMap;
use std::io;
use std::os::unix::fs::PermissionsExt;
use std::path::{Path, PathBuf};
use std::sync::{Arc, Mutex};
use tokio::net::UnixListener;
use tokio::task::{JoinHandle, JoinSet};
use tracing::warn;
use crate::api::BuildCacheCounts;
use crate::config::BuildCacheWriters;
use crate::github::{PullRequest, RunTrigger};
use crate::sync::lock;
pub const MOUNT: &str = "/run/runnerd-build-cache";
pub const SOCKET: &str = "cache.sock";
const ENTRIES: &str = "entries";
const UPLOADS: &str = "uploads";
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum Writes {
Undecided,
Allowed,
Refused,
}
const WRITING_EVENTS: [&str; 3] = ["push", "schedule", "workflow_dispatch"];
pub const PULL_REQUEST_EVENT: &str = "pull_request";
pub fn writes(
runner: &str,
job_runner: Option<&str>,
repository: &str,
trigger: &RunTrigger,
branches: &[String],
writers: BuildCacheWriters,
) -> Writes {
match job_runner {
None => Writes::Undecided,
Some(placed) if placed != runner => Writes::Refused,
Some(_) if !same_repository(trigger.head_repository.as_deref(), repository) => {
Writes::Refused
}
Some(_) => {
let Some(branch) = &trigger.head_branch else {
return Writes::Refused;
};
let event = trigger.event.as_str();
let branch_run = WRITING_EVENTS.contains(&event) && branches.contains(branch);
let pull_request = writers == BuildCacheWriters::PullRequests
&& event == PULL_REQUEST_EVENT
&& trigger.by_user
&& !trigger.pull_requests.is_empty();
if branch_run || pull_request {
Writes::Allowed
} else {
Writes::Refused
}
}
}
}
pub fn opened_by_user(pull: &PullRequest, repository: &str, trigger: &RunTrigger) -> bool {
pull.by_user
&& same_repository(pull.head_repository.as_deref(), repository)
&& trigger.head_branch.as_ref() == Some(&pull.head_branch)
}
fn same_repository(head: Option<&str>, repository: &str) -> bool {
head.is_some_and(|head| head.eq_ignore_ascii_case(repository))
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Size {
pub bytes: u64,
pub entries: u64,
}
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub struct Trimmed {
pub before: Size,
pub after: Size,
}
#[derive(Debug, Clone)]
pub struct BuildCache {
inner: Arc<Inner>,
}
#[derive(Debug)]
struct Inner {
root: PathBuf,
uploads: PathBuf,
sockets: PathBuf,
state: Mutex<State>,
}
#[derive(Debug, Default)]
struct State {
max_size: u64,
entries: HashMap<PathBuf, Entry>,
bytes: u64,
runners: HashMap<String, Runner>,
next_upload: u64,
uploading: u64,
}
#[derive(Debug, Clone, Copy)]
struct Entry {
bytes: u64,
used: u64,
}
#[derive(Debug)]
struct Runner {
repository: String,
writes: Writes,
counts: BuildCacheCounts,
held: Vec<Upload>,
uploading: u64,
dir: PathBuf,
serving: JoinHandle<()>,
}
#[derive(Debug)]
struct Upload {
file: PathBuf,
entry: PathBuf,
bytes: u64,
}
impl BuildCache {
pub fn open(directory: &Path, sockets: &Path, max_size: u64) -> io::Result<BuildCache> {
std::fs::create_dir_all(directory)?;
let (root, uploads) = (directory.join(ENTRIES), directory.join(UPLOADS));
match std::fs::remove_dir_all(&uploads) {
Err(error) if error.kind() != io::ErrorKind::NotFound => return Err(error),
_ => {}
}
private_dir(&root)?;
private_dir(&uploads)?;
private_dir(sockets)?;
let mut entries = HashMap::new();
count_entries(&root, &root, &mut entries)?;
let bytes = entries.values().map(|entry: &Entry| entry.bytes).sum();
Ok(BuildCache {
inner: Arc::new(Inner {
root,
uploads,
sockets: sockets.to_owned(),
state: Mutex::new(State {
max_size,
entries,
bytes,
..State::default()
}),
}),
})
}
pub fn directory(&self) -> &Path {
&self.inner.root
}
pub fn set_max_size(&self, max_size: u64) {
lock(&self.inner.state).max_size = max_size;
}
pub fn attach(&self, runner: &str, repository: &str) -> io::Result<PathBuf> {
if !valid_repository(repository) {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("{repository} is not a repository name"),
));
}
let dir = self.inner.sockets.join(socket_dir(runner));
std::fs::create_dir_all(&dir)?;
std::fs::set_permissions(&dir, std::fs::Permissions::from_mode(0o755))?;
let socket = dir.join(SOCKET);
match std::fs::remove_file(&socket) {
Err(error) if error.kind() != io::ErrorKind::NotFound => return Err(error),
_ => {}
}
let listener = UnixListener::bind(&socket)?;
std::fs::set_permissions(&socket, std::fs::Permissions::from_mode(0o666))?;
let serving = tokio::spawn(accept(self.inner.clone(), runner.to_owned(), listener));
let old = lock(&self.inner.state).runners.insert(
runner.to_owned(),
Runner {
repository: repository.to_owned(),
writes: Writes::Undecided,
counts: BuildCacheCounts::default(),
held: Vec::new(),
uploading: 0,
dir: dir.clone(),
serving,
},
);
if let Some(old) = old {
old.serving.abort();
self.inner.remove_uploads(None, &old.held);
}
Ok(dir)
}
pub fn forget_others(&self, keep: &[&str]) {
let keep: Vec<String> = keep.iter().map(|runner| socket_dir(runner)).collect();
let Ok(dirs) = std::fs::read_dir(&self.inner.sockets) else {
return;
};
for dir in dirs.flatten() {
if !keep
.iter()
.any(|name| *name == dir.file_name().to_string_lossy())
{
let _ = std::fs::remove_dir_all(dir.path());
}
}
}
pub fn undecided(&self, runner: &str) -> bool {
lock(&self.inner.state)
.runners
.get(runner)
.is_some_and(|runner| runner.writes == Writes::Undecided)
}
pub fn decide(&self, runner: &str, writes: Writes) {
let held = {
let mut state = lock(&self.inner.state);
let Some(open) = state.runners.get_mut(runner) else {
return;
};
if open.writes != Writes::Undecided || writes == Writes::Undecided {
return;
}
open.writes = writes;
std::mem::take(&mut open.held)
};
match writes {
Writes::Allowed => {
for upload in held {
self.inner.store(runner, upload);
}
}
Writes::Refused | Writes::Undecided => {
if let Some(open) = lock(&self.inner.state).runners.get_mut(runner) {
open.counts.refused += held.len() as u64;
}
self.inner.remove_uploads(Some(runner), &held);
}
}
}
pub fn close(&self, runner: &str) -> bool {
let dir = {
let state = lock(&self.inner.state);
let Some(open) = state.runners.get(runner) else {
return false;
};
if open.writes != Writes::Undecided || open.held.is_empty() {
return false;
}
open.serving.abort();
open.dir.clone()
};
if let Err(error) = std::fs::remove_dir_all(&dir)
&& error.kind() != io::ErrorKind::NotFound
{
warn!("cannot remove {}: {error}", dir.display());
}
true
}
pub fn detach(&self, runner: &str) -> Option<BuildCacheCounts> {
let Some(mut open) = lock(&self.inner.state).runners.remove(runner) else {
let dir = self.inner.sockets.join(socket_dir(runner));
if let Err(error) = std::fs::remove_dir_all(&dir)
&& error.kind() != io::ErrorKind::NotFound
{
warn!("cannot remove {}: {error}", dir.display());
}
return None;
};
open.serving.abort();
self.inner.remove_uploads(None, &open.held);
open.counts.refused += open.held.len() as u64;
open.counts.writer = open.writes == Writes::Allowed;
if let Err(error) = std::fs::remove_dir_all(&open.dir)
&& error.kind() != io::ErrorKind::NotFound
{
warn!("cannot remove {}: {error}", open.dir.display());
}
Some(open.counts)
}
pub fn size(&self) -> Size {
let state = lock(&self.inner.state);
Size {
bytes: state.bytes,
entries: state.entries.len() as u64,
}
}
pub fn trim(&self) -> Trimmed {
let limit = lock(&self.inner.state).max_size;
self.trim_to(limit)
}
pub fn clear(&self) -> Trimmed {
self.trim_to(0)
}
fn trim_to(&self, limit: u64) -> Trimmed {
let (before, victims) = {
let mut state = lock(&self.inner.state);
let before = Size {
bytes: state.bytes,
entries: state.entries.len() as u64,
};
let mut by_use: Vec<(PathBuf, Entry)> = state
.entries
.iter()
.map(|(path, entry)| (path.clone(), *entry))
.collect();
by_use.sort_by_key(|(path, entry)| (entry.used, path.clone()));
let mut victims = Vec::new();
for (path, entry) in by_use {
if state.bytes <= limit && (limit > 0 || state.entries.is_empty()) {
break;
}
state.entries.remove(&path);
state.bytes -= entry.bytes;
victims.push(path);
}
(before, victims)
};
for path in victims {
if let Err(error) = std::fs::remove_file(self.inner.root.join(&path))
&& error.kind() != io::ErrorKind::NotFound
{
warn!(
"cannot remove build cache entry {}: {error}",
path.display()
);
}
}
Trimmed {
before,
after: self.size(),
}
}
}
impl Inner {
fn store(&self, runner: &str, upload: Upload) {
let mut state = lock(&self.state);
let replaced = state
.entries
.get(&upload.entry)
.map_or(0, |entry| entry.bytes);
let fits = state.bytes - replaced + upload.bytes <= state.max_size.saturating_mul(2);
let path = self.root.join(&upload.entry);
let stored = fits
&& path
.parent()
.is_none_or(|parent| std::fs::create_dir_all(parent).is_ok())
&& std::fs::rename(&upload.file, &path).is_ok();
Self::give_back(&mut state, Some(runner), upload.bytes);
if stored {
let entry = Entry {
bytes: upload.bytes,
used: crate::time::now_ms(),
};
if let Some(old) = state.entries.insert(upload.entry, entry) {
state.bytes -= old.bytes;
}
state.bytes += upload.bytes;
} else {
let _ = std::fs::remove_file(&upload.file);
}
if let Some(open) = state.runners.get_mut(runner) {
if stored {
open.counts.written_bytes += upload.bytes;
} else {
open.counts.refused += 1;
}
}
}
fn reserve(&self, runner: &str, bytes: u64) -> bool {
let mut state = lock(&self.state);
let (total, max_size) = (state.uploading, state.max_size);
let Some(open) = state.runners.get_mut(runner) else {
return false;
};
let fits = total + bytes <= max_size
&& (open.writes != Writes::Undecided || open.uploading + bytes <= max_size / 4);
if fits {
open.uploading += bytes;
state.uploading += bytes;
}
fits
}
fn release(&self, runner: Option<&str>, bytes: u64) {
Self::give_back(&mut lock(&self.state), runner, bytes);
}
fn give_back(state: &mut State, runner: Option<&str>, bytes: u64) {
state.uploading = state.uploading.saturating_sub(bytes);
if let Some(open) = runner.and_then(|runner| state.runners.get_mut(runner)) {
open.uploading = open.uploading.saturating_sub(bytes);
}
}
fn remove_uploads(&self, runner: Option<&str>, uploads: &[Upload]) {
for upload in uploads {
let _ = std::fs::remove_file(&upload.file);
self.release(runner, upload.bytes);
}
}
fn touch(&self, entry: &Path) {
if let Some(used) = lock(&self.state).entries.get_mut(entry) {
used.used = crate::time::now_ms();
}
let path = self.root.join(entry);
tokio::task::spawn_blocking(move || {
if let Ok(file) = std::fs::File::open(&path) {
let _ = file.set_modified(std::time::SystemTime::now());
}
});
}
}
async fn accept(inner: Arc<Inner>, runner: String, listener: UnixListener) {
let mut connections = JoinSet::new();
loop {
tokio::select! {
accepted = listener.accept() => match accepted {
Ok((stream, _)) => {
connections.spawn(serve::connection(inner.clone(), runner.clone(), stream));
}
Err(error) => {
warn!("{runner}: the build cache's socket failed: {error}");
return;
}
},
Some(_) = connections.join_next() => {}
}
}
}
fn socket_dir(runner: &str) -> String {
let hash = runner
.bytes()
.fold(0xcbf2_9ce4_8422_2325_u64, |hash, byte| {
(hash ^ u64::from(byte)).wrapping_mul(0x0000_0100_0000_01b3)
});
format!("{hash:016x}")
}
fn valid_segment(name: &str) -> bool {
!name.is_empty()
&& name.len() <= 200
&& name != "."
&& name != ".."
&& name
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'.' | b'_' | b'-'))
}
fn valid_repository(repository: &str) -> bool {
repository.split_once('/').is_some_and(|(owner, name)| {
valid_segment(owner) && valid_segment(name) && !owner.starts_with('.')
})
}
fn private_dir(path: &Path) -> io::Result<()> {
std::fs::create_dir_all(path)?;
std::fs::set_permissions(path, std::fs::Permissions::from_mode(0o700))
}
fn count_entries(root: &Path, dir: &Path, entries: &mut HashMap<PathBuf, Entry>) -> io::Result<()> {
for item in std::fs::read_dir(dir)? {
let item = item?;
let path = item.path();
let metadata = item.metadata()?;
if metadata.is_dir() {
count_entries(root, &path, entries)?;
} else if metadata.is_file()
&& let Ok(relative) = path.strip_prefix(root)
{
entries.insert(
relative.to_owned(),
Entry {
bytes: metadata.len(),
used: metadata.modified().map_or(0, crate::time::ms),
},
);
}
}
Ok(())
}
#[cfg(test)]
mod tests;
#[cfg(test)]
#[allow(
clippy::unwrap_used,
reason = "a test helper: a socket that does not answer fails the test"
)]
pub(crate) async fn request(dir: &Path, method: &str, path: &str, body: &[u8]) -> (u16, Vec<u8>) {
use tokio::io::{AsyncReadExt, AsyncWriteExt};
let mut stream = tokio::net::UnixStream::connect(dir.join(SOCKET))
.await
.unwrap();
let head = format!(
"{method} {path} HTTP/1.1\r\nHost: 127.0.0.1\r\nAuthorization: Basic dG9rZW46c2VjcmV0\r\n\
Content-Length: {}\r\nConnection: close\r\n\r\n",
body.len()
);
stream.write_all(head.as_bytes()).await.unwrap();
stream.write_all(body).await.unwrap();
let mut response = Vec::new();
stream.read_to_end(&mut response).await.unwrap();
let split = response
.windows(4)
.position(|window| window == b"\r\n\r\n")
.unwrap();
let status = String::from_utf8_lossy(&response[9..12]).parse().unwrap();
(status, response[split + 4..].to_vec())
}