feat: compute shared_buffers and mount /dev/shm for restore database (#11)

Author: passcod

src/controllers/restore/builders.rs

@@ -18,6 +18,7 @@ use super::restore_owner_reference;
 use crate::{
 	controllers::{env_from_secret, env_from_secret_optional, kopia_writable_env, overlay},
 	error::{Error, Result},
+	quantity::compute_shm_and_shared_buffers,
 	types::*,
 };
 
@@ -411,6 +412,7 @@ pub fn build_deployment(
 	replica: &PostgresPhysicalReplica,
 ) -> Result<Deployment> {
 	let pvc_name = format!("{name}-data");
+	let (shm_size, shared_buffers_mb) = compute_shm_and_shared_buffers(&replica.spec.resources);
 	let creds_secret = SecretReference {
 		name: Some(format!("{}-creds", restore.spec.replica.name)),
 		namespace: Some(namespace.to_string()),
@@ -506,7 +508,7 @@ listen_addresses = '*'
 port = 5432
 max_connections = 100
 max_prepared_transactions = 16
-shared_buffers = 128MB
+shared_buffers = {shared_buffers_mb}MB
 dynamic_shared_memory_type = posix
 log_timezone = 'UTC'
 datestyle = 'iso, mdy'
@@ -526,6 +528,7 @@ sed -i \
   -e '/^[[:space:]]*ident_file[[:space:]]*=/d' \
   -e '/^[[:space:]]*data_directory[[:space:]]*=/d' \
   -e '/^[[:space:]]*dynamic_shared_memory_type[[:space:]]*=/d' \
+  -e '/^[[:space:]]*shared_buffers[[:space:]]*=/d' \
   -e '/^[[:space:]]*log_destination[[:space:]]*=/d' \
   -e '/^[[:space:]]*logging_collector[[:space:]]*=/d' \
   -e '/^[[:space:]]*log_directory[[:space:]]*=/d' \
@@ -546,6 +549,7 @@ echo "Configuring stderr logging..."
 echo "log_destination = 'stderr'" >> "$PGDATA/postgresql.conf"
 echo "password_encryption = 'scram-sha-256'" >> "$PGDATA/postgresql.conf"
 echo "logging_collector = off" >> "$PGDATA/postgresql.conf"
+echo "shared_buffers = {shared_buffers_mb}MB" >> "$PGDATA/postgresql.conf"
 
 PG_MAJOR=$(cat "$PGDATA/PG_VERSION")
 
@@ -908,6 +912,11 @@ exec postgres -D /pgdata/pgdata ${PGRO_LOG_LEVEL:+-c log_min_messages=$PGRO_LOG_
 								mount_path: "/usr/lib/locale".to_string(),
 								..Default::default()
 							},
+							VolumeMount {
+								name: "dshm".to_string(),
+								mount_path: "/dev/shm".to_string(),
+								..Default::default()
+							},
 						]),
 						readiness_probe: Some(Probe {
 							exec: Some(ExecAction {
@@ -960,6 +969,16 @@ exec postgres -D /pgdata/pgdata ${PGRO_LOG_LEVEL:+-c log_min_messages=$PGRO_LOG_
 							),
 							..Default::default()
 						},
+						Volume {
+							name: "dshm".to_string(),
+							empty_dir: Some(
+								k8s_openapi::api::core::v1::EmptyDirVolumeSource {
+										medium: Some("Memory".to_string()),
+										size_limit: Some(shm_size),
+									},
+							),
+							..Default::default()
+						},
 					]),
 					affinity: replica.spec.affinity.clone(),
 					tolerations: Some(replica.spec.tolerations.clone()),

src/controllers/restore/tests.rs

@@ -1,7 +1,9 @@
+use std::collections::BTreeMap;
+
 use jiff::Span;
 use k8s_openapi::api::core::v1::{
 	Affinity, LocalObjectReference, NodeAffinity, NodeSelector, NodeSelectorRequirement,
-	NodeSelectorTerm, SecretReference,
+	NodeSelectorTerm, ResourceRequirements, SecretReference,
 };
 use k8s_openapi::apimachinery::pkg::api::resource::Quantity;
 
@@ -172,3 +174,169 @@ fn version_detect_job_has_ttl_seconds_after_finished() {
 		.expect("version-detect job must set ttlSecondsAfterFinished");
 	assert!(ttl > 0, "ttlSecondsAfterFinished must be positive");
 }
+
+#[test]
+fn deployment_has_dshm_volume_with_default_resources() {
+	let (mut restore, replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let dshm_vol = pod_spec
+		.volumes
+		.as_ref()
+		.unwrap()
+		.iter()
+		.find(|v| v.name == "dshm")
+		.expect("dshm volume must exist");
+
+	let empty_dir = dshm_vol.empty_dir.as_ref().expect("dshm must be emptyDir");
+	assert_eq!(empty_dir.medium.as_deref(), Some("Memory"));
+	// Default 1Gi request: min(512Mi, ceil(36% of 1Gi)) = min(512Mi, 369Mi) = 369Mi
+	assert_eq!(empty_dir.size_limit.as_ref().unwrap().0, "369Mi");
+}
+
+#[test]
+fn deployment_has_dshm_volume_with_custom_resources() {
+	let (mut restore, mut replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+	replica.spec.resources = Some(ResourceRequirements {
+		requests: Some(BTreeMap::from([(
+			"memory".to_string(),
+			Quantity("4Gi".to_string()),
+		)])),
+		limits: Some(BTreeMap::from([(
+			"memory".to_string(),
+			Quantity("8Gi".to_string()),
+		)])),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let dshm_vol = pod_spec
+		.volumes
+		.as_ref()
+		.unwrap()
+		.iter()
+		.find(|v| v.name == "dshm")
+		.expect("dshm volume must exist");
+
+	let empty_dir = dshm_vol.empty_dir.as_ref().expect("dshm must be emptyDir");
+	assert_eq!(empty_dir.medium.as_deref(), Some("Memory"));
+	// 4Gi request, 8Gi limit: min(4Gi/2, ceil(36% of 8Gi)) = min(2048Mi, 2950Mi) = 2048Mi
+	assert_eq!(empty_dir.size_limit.as_ref().unwrap().0, "2048Mi");
+}
+
+#[test]
+fn deployment_mounts_dshm_on_postgres_and_setup_auth() {
+	let (mut restore, replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let postgres = &pod_spec.containers[0];
+	assert_eq!(postgres.name, "postgres");
+	let pg_mounts = postgres.volume_mounts.as_ref().unwrap();
+	let pg_shm = pg_mounts
+		.iter()
+		.find(|m| m.mount_path == "/dev/shm")
+		.expect("postgres container must mount /dev/shm");
+	assert_eq!(pg_shm.name, "dshm");
+
+	let setup_auth = pod_spec
+		.init_containers
+		.as_ref()
+		.unwrap()
+		.iter()
+		.find(|c| c.name == "setup-auth")
+		.expect("setup-auth init container must exist");
+	let sa_mounts = setup_auth.volume_mounts.as_ref().unwrap();
+	assert!(
+		!sa_mounts.iter().any(|m| m.mount_path == "/dev/shm"),
+		"setup-auth must NOT mount /dev/shm (it never uses shared memory)"
+	);
+}
+
+#[test]
+fn deployment_init_script_sets_shared_buffers() {
+	let (mut restore, replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let setup_auth = pod_spec
+		.init_containers
+		.as_ref()
+		.unwrap()
+		.iter()
+		.find(|c| c.name == "setup-auth")
+		.expect("setup-auth init container must exist");
+	let script = &setup_auth.args.as_ref().unwrap()[0];
+
+	// Default 1Gi: SHM=369Mi, shared_buffers = floor(70% of 369) = 258MB
+	assert!(
+		script.contains("shared_buffers = 258MB"),
+		"init script must set shared_buffers to computed value, got script containing: {}",
+		script
+			.lines()
+			.find(|l| l.contains("shared_buffers"))
+			.unwrap_or("<not found>")
+	);
+
+	// The sed block must strip shared_buffers from source config
+	assert!(
+		script.contains("shared_buffers[[:space:]]*="),
+		"sed must strip shared_buffers from source config"
+	);
+}
+
+#[test]
+fn deployment_shared_buffers_with_custom_resources() {
+	let (mut restore, mut replica) = test_restore_and_replica();
+	restore.status = Some(PostgresPhysicalRestoreStatus {
+		postgres_version: Some("16".to_string()),
+		..Default::default()
+	});
+	replica.spec.resources = Some(ResourceRequirements {
+		requests: Some(BTreeMap::from([(
+			"memory".to_string(),
+			Quantity("2Gi".to_string()),
+		)])),
+		..Default::default()
+	});
+
+	let deploy = build_deployment(&restore, "test-restore", "default", &replica).unwrap();
+	let pod_spec = deploy.spec.unwrap().template.spec.unwrap();
+
+	let setup_auth = pod_spec
+		.init_containers
+		.as_ref()
+		.unwrap()
+		.iter()
+		.find(|c| c.name == "setup-auth")
+		.expect("setup-auth init container must exist");
+	let script = &setup_auth.args.as_ref().unwrap()[0];
+
+	// 2Gi request: SHM=738Mi, shared_buffers = floor(70% of 738) = 516MB
+	assert!(
+		script.contains("shared_buffers = 516MB"),
+		"init script must set shared_buffers for 2Gi request"
+	);
+}