diff --git a/compute/vm-image-spec.yaml b/compute/vm-image-spec.yaml index 0af44745e5..50fcd62e4f 100644 --- a/compute/vm-image-spec.yaml +++ b/compute/vm-image-spec.yaml @@ -11,6 +11,10 @@ commands: user: root sysvInitAction: sysinit shell: 'chmod 711 /neonvm/bin/resize-swap' + - name: chmod-set-disk-quota + user: root + sysvInitAction: sysinit + shell: 'chmod 711 /neonvm/bin/set-disk-quota' - name: pgbouncer user: postgres sysvInitAction: respawn @@ -30,11 +34,12 @@ commands: shutdownHook: | su -p postgres --session-command '/usr/local/bin/pg_ctl stop -D /var/db/postgres/compute/pgdata -m fast --wait -t 10' files: - - filename: compute_ctl-resize-swap + - filename: compute_ctl-sudoers content: | # Allow postgres user (which is what compute_ctl runs as) to run /neonvm/bin/resize-swap - # as root without requiring entering a password (NOPASSWD), regardless of hostname (ALL) - postgres ALL=(root) NOPASSWD: /neonvm/bin/resize-swap + # and /neonvm/bin/set-disk-quota as root without requiring entering a password (NOPASSWD), + # regardless of hostname (ALL) + postgres ALL=(root) NOPASSWD: /neonvm/bin/resize-swap, /neonvm/bin/set-disk-quota - filename: cgconfig.conf content: | # Configuration for cgroups in VM compute nodes @@ -100,7 +105,7 @@ merge: | && apt install --no-install-recommends -y \ sudo \ && rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* - COPY compute_ctl-resize-swap /etc/sudoers.d/compute_ctl-resize-swap + COPY compute_ctl-sudoers /etc/sudoers.d/compute_ctl-sudoers COPY cgconfig.conf /etc/cgconfig.conf diff --git a/compute_tools/src/bin/compute_ctl.rs b/compute_tools/src/bin/compute_ctl.rs index 9499a7186e..b10638c454 100644 --- a/compute_tools/src/bin/compute_ctl.rs +++ b/compute_tools/src/bin/compute_ctl.rs @@ -44,6 +44,7 @@ use std::{thread, time::Duration}; use anyhow::{Context, Result}; use chrono::Utc; use clap::Arg; +use compute_tools::disk_quota::set_disk_quota; use compute_tools::lsn_lease::launch_lsn_lease_bg_task_for_static; use signal_hook::consts::{SIGQUIT, SIGTERM}; use signal_hook::{consts::SIGINT, iterator::Signals}; @@ -151,6 +152,7 @@ fn process_cli(matches: &clap::ArgMatches) -> Result { let spec_json = matches.get_one::("spec"); let spec_path = matches.get_one::("spec-path"); let resize_swap_on_bind = matches.get_flag("resize-swap-on-bind"); + let set_disk_quota_for_fs = matches.get_one::("set-disk-quota-for-fs"); Ok(ProcessCliResult { connstr, @@ -161,6 +163,7 @@ fn process_cli(matches: &clap::ArgMatches) -> Result { spec_json, spec_path, resize_swap_on_bind, + set_disk_quota_for_fs, }) } @@ -173,6 +176,7 @@ struct ProcessCliResult<'clap> { spec_json: Option<&'clap String>, spec_path: Option<&'clap String>, resize_swap_on_bind: bool, + set_disk_quota_for_fs: Option<&'clap String>, } fn startup_context_from_env() -> Option { @@ -293,6 +297,7 @@ fn wait_spec( pgbin, ext_remote_storage, resize_swap_on_bind, + set_disk_quota_for_fs, http_port, .. }: ProcessCliResult, @@ -373,6 +378,7 @@ fn wait_spec( compute, http_port, resize_swap_on_bind, + set_disk_quota_for_fs: set_disk_quota_for_fs.cloned(), }) } @@ -381,6 +387,7 @@ struct WaitSpecResult { // passed through from ProcessCliResult http_port: u16, resize_swap_on_bind: bool, + set_disk_quota_for_fs: Option, } fn start_postgres( @@ -390,6 +397,7 @@ fn start_postgres( compute, http_port, resize_swap_on_bind, + set_disk_quota_for_fs, }: WaitSpecResult, ) -> Result<(Option, StartPostgresResult)> { // We got all we need, update the state. @@ -403,6 +411,7 @@ fn start_postgres( ); // before we release the mutex, fetch the swap size (if any) for later. let swap_size_bytes = state.pspec.as_ref().unwrap().spec.swap_size_bytes; + let disk_quota_bytes = state.pspec.as_ref().unwrap().spec.disk_quota_bytes; drop(state); // Launch remaining service threads @@ -422,8 +431,8 @@ fn start_postgres( // OOM-killed during startup because swap wasn't available yet. match resize_swap(size_bytes) { Ok(()) => { - let size_gib = size_bytes as f32 / (1 << 20) as f32; // just for more coherent display. - info!(%size_bytes, %size_gib, "resized swap"); + let size_mib = size_bytes as f32 / (1 << 20) as f32; // just for more coherent display. + info!(%size_bytes, %size_mib, "resized swap"); } Err(err) => { let err = err.context("failed to resize swap"); @@ -432,10 +441,29 @@ fn start_postgres( // Mark compute startup as failed; don't try to start postgres, and report this // error to the control plane when it next asks. prestartup_failed = true; - let mut state = compute.state.lock().unwrap(); - state.error = Some(format!("{err:?}")); - state.status = ComputeStatus::Failed; - compute.state_changed.notify_all(); + compute.set_failed_status(err); + delay_exit = true; + } + } + } + + // Set disk quota if the compute spec says so + if let (Some(disk_quota_bytes), Some(disk_quota_fs_mountpoint)) = + (disk_quota_bytes, set_disk_quota_for_fs) + { + match set_disk_quota(disk_quota_bytes, &disk_quota_fs_mountpoint) { + Ok(()) => { + let size_mib = disk_quota_bytes as f32 / (1 << 20) as f32; // just for more coherent display. + info!(%disk_quota_bytes, %size_mib, "set disk quota"); + } + Err(err) => { + let err = err.context("failed to set disk quota"); + error!("{err:#}"); + + // Mark compute startup as failed; don't try to start postgres, and report this + // error to the control plane when it next asks. + prestartup_failed = true; + compute.set_failed_status(err); delay_exit = true; } } @@ -450,16 +478,7 @@ fn start_postgres( Ok(pg) => Some(pg), Err(err) => { error!("could not start the compute node: {:#}", err); - let mut state = compute.state.lock().unwrap(); - state.error = Some(format!("{:?}", err)); - state.status = ComputeStatus::Failed; - // Notify others that Postgres failed to start. In case of configuring the - // empty compute, it's likely that API handler is still waiting for compute - // state change. With this we will notify it that compute is in Failed state, - // so control plane will know about it earlier and record proper error instead - // of timeout. - compute.state_changed.notify_all(); - drop(state); // unlock + compute.set_failed_status(err); delay_exit = true; None } @@ -750,6 +769,11 @@ fn cli() -> clap::Command { .long("resize-swap-on-bind") .action(clap::ArgAction::SetTrue), ) + .arg( + Arg::new("set-disk-quota-for-fs") + .long("set-disk-quota-for-fs") + .value_name("SET_DISK_QUOTA_FOR_FS") + ) } /// When compute_ctl is killed, send also termination signal to sync-safekeepers diff --git a/compute_tools/src/compute.rs b/compute_tools/src/compute.rs index a6b7633eda..147eb2a161 100644 --- a/compute_tools/src/compute.rs +++ b/compute_tools/src/compute.rs @@ -306,6 +306,13 @@ impl ComputeNode { self.state_changed.notify_all(); } + pub fn set_failed_status(&self, err: anyhow::Error) { + let mut state = self.state.lock().unwrap(); + state.error = Some(format!("{err:?}")); + state.status = ComputeStatus::Failed; + self.state_changed.notify_all(); + } + pub fn get_status(&self) -> ComputeStatus { self.state.lock().unwrap().status } diff --git a/compute_tools/src/disk_quota.rs b/compute_tools/src/disk_quota.rs new file mode 100644 index 0000000000..e838c5b9fd --- /dev/null +++ b/compute_tools/src/disk_quota.rs @@ -0,0 +1,25 @@ +use anyhow::Context; + +pub const DISK_QUOTA_BIN: &str = "/neonvm/bin/set-disk-quota"; + +/// If size_bytes is 0, it disables the quota. Otherwise, it sets filesystem quota to size_bytes. +/// `fs_mountpoint` should point to the mountpoint of the filesystem where the quota should be set. +pub fn set_disk_quota(size_bytes: u64, fs_mountpoint: &str) -> anyhow::Result<()> { + let size_kb = size_bytes / 1024; + // run `/neonvm/bin/set-disk-quota {size_kb} {mountpoint}` + let child_result = std::process::Command::new("/usr/bin/sudo") + .arg(DISK_QUOTA_BIN) + .arg(size_kb.to_string()) + .arg(fs_mountpoint) + .spawn(); + + child_result + .context("spawn() failed") + .and_then(|mut child| child.wait().context("wait() failed")) + .and_then(|status| match status.success() { + true => Ok(()), + false => Err(anyhow::anyhow!("process exited with {status}")), + }) + // wrap any prior error with the overall context that we couldn't run the command + .with_context(|| format!("could not run `/usr/bin/sudo {DISK_QUOTA_BIN}`")) +} diff --git a/compute_tools/src/lib.rs b/compute_tools/src/lib.rs index c402d63305..c5b4ca632c 100644 --- a/compute_tools/src/lib.rs +++ b/compute_tools/src/lib.rs @@ -10,6 +10,7 @@ pub mod http; pub mod logger; pub mod catalog; pub mod compute; +pub mod disk_quota; pub mod extension_server; pub mod lsn_lease; mod migration; diff --git a/control_plane/src/endpoint.rs b/control_plane/src/endpoint.rs index 7554a03a68..18f396b886 100644 --- a/control_plane/src/endpoint.rs +++ b/control_plane/src/endpoint.rs @@ -561,6 +561,7 @@ impl Endpoint { operation_uuid: None, features: self.features.clone(), swap_size_bytes: None, + disk_quota_bytes: None, cluster: Cluster { cluster_id: None, // project ID: not used name: None, // project name: not used diff --git a/libs/compute_api/src/spec.rs b/libs/compute_api/src/spec.rs index 525a1572ff..83515a00a0 100644 --- a/libs/compute_api/src/spec.rs +++ b/libs/compute_api/src/spec.rs @@ -50,6 +50,16 @@ pub struct ComputeSpec { #[serde(default)] pub swap_size_bytes: Option, + /// If compute_ctl was passed `--set-disk-quota-for-fs`, a value of `Some(_)` instructs + /// compute_ctl to run `/neonvm/bin/set-disk-quota` with the given size and fs, when the + /// spec is first received. + /// + /// Both this field and `--set-disk-quota-for-fs` are required, so that the control plane's + /// spec generation doesn't need to be aware of the actual compute it's running on, while + /// guaranteeing gradual rollout of disk quota. + #[serde(default)] + pub disk_quota_bytes: Option, + /// Expected cluster state at the end of transition process. pub cluster: Cluster, pub delta_operations: Option>,