diff options
| author | soryu <soryu@soryu.co> | 2026-01-22 22:32:46 +0000 |
|---|---|---|
| committer | soryu <soryu@soryu.co> | 2026-01-23 01:03:04 +0000 |
| commit | 1ed362424dafec690f919154f5116471951cda9c (patch) | |
| tree | 19c7ca9231887394a791223fe32a8ad335a687a8 /makima/src/daemon/storage/patch.rs | |
| parent | 265f8cf14fec9d7116d09af49e4b48b357faceda (diff) | |
| download | soryu-1ed362424dafec690f919154f5116471951cda9c.tar.gz soryu-1ed362424dafec690f919154f5116471951cda9c.zip | |
Add patch checkpointing
Diffstat (limited to 'makima/src/daemon/storage/patch.rs')
| -rw-r--r-- | makima/src/daemon/storage/patch.rs | 293 |
1 files changed, 293 insertions, 0 deletions
diff --git a/makima/src/daemon/storage/patch.rs b/makima/src/daemon/storage/patch.rs new file mode 100644 index 0000000..45624b5 --- /dev/null +++ b/makima/src/daemon/storage/patch.rs @@ -0,0 +1,293 @@ +//! Git patch creation and application for checkpoint recovery. + +use flate2::read::GzDecoder; +use flate2::write::GzEncoder; +use flate2::Compression; +use std::io::{Read, Write}; +use std::path::Path; +use thiserror::Error; +use tokio::process::Command; + +/// Errors that can occur during patch operations. +#[derive(Error, Debug)] +pub enum PatchError { + #[error("Git command failed: {0}")] + GitCommand(String), + + #[error("Compression error: {0}")] + Compression(#[from] std::io::Error), + + #[error("Patch too large: {size} bytes (max: {max} bytes)")] + TooLarge { size: usize, max: usize }, + + #[error("Empty patch (no changes)")] + EmptyPatch, + + #[error("Failed to apply patch: {0}")] + ApplyFailed(String), +} + +/// Create a compressed git diff from worktree changes. +/// +/// Generates a diff between `base_sha` and HEAD, then compresses it with gzip. +/// Returns the compressed patch bytes and the number of files changed. +pub async fn create_patch( + worktree_path: &Path, + base_sha: &str, +) -> Result<(Vec<u8>, usize), PatchError> { + // Get the diff between base commit and HEAD + let output = Command::new("git") + .current_dir(worktree_path) + .args(["diff", base_sha, "HEAD", "--binary"]) + .output() + .await + .map_err(|e| PatchError::GitCommand(format!("Failed to run git diff: {}", e)))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(PatchError::GitCommand(format!("git diff failed: {}", stderr))); + } + + let diff_data = output.stdout; + if diff_data.is_empty() { + return Err(PatchError::EmptyPatch); + } + + // Count files changed + let files_output = Command::new("git") + .current_dir(worktree_path) + .args(["diff", base_sha, "HEAD", "--name-only"]) + .output() + .await + .map_err(|e| PatchError::GitCommand(format!("Failed to count files: {}", e)))?; + + let files_count = if files_output.status.success() { + String::from_utf8_lossy(&files_output.stdout) + .lines() + .filter(|l| !l.is_empty()) + .count() + } else { + 0 + }; + + // Compress with gzip + let mut encoder = GzEncoder::new(Vec::new(), Compression::default()); + encoder.write_all(&diff_data)?; + let compressed = encoder.finish()?; + + Ok((compressed, files_count)) +} + +/// Apply a compressed patch to restore worktree state. +/// +/// The worktree should already be checked out at `base_sha` before calling this. +pub async fn apply_patch(worktree_path: &Path, patch_data: &[u8]) -> Result<(), PatchError> { + // Decompress gzip + let mut decoder = GzDecoder::new(patch_data); + let mut decompressed = Vec::new(); + decoder.read_to_end(&mut decompressed)?; + + if decompressed.is_empty() { + return Err(PatchError::EmptyPatch); + } + + // Apply the patch using git apply + let mut child = Command::new("git") + .current_dir(worktree_path) + .args(["apply", "--binary", "-"]) + .stdin(std::process::Stdio::piped()) + .stdout(std::process::Stdio::piped()) + .stderr(std::process::Stdio::piped()) + .spawn() + .map_err(|e| PatchError::GitCommand(format!("Failed to spawn git apply: {}", e)))?; + + // Write patch to stdin + if let Some(mut stdin) = child.stdin.take() { + use tokio::io::AsyncWriteExt; + stdin.write_all(&decompressed).await?; + drop(stdin); // Close stdin to signal EOF + } + + let output = child + .wait_with_output() + .await + .map_err(|e| PatchError::GitCommand(format!("Failed to wait for git apply: {}", e)))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(PatchError::ApplyFailed(stderr.to_string())); + } + + Ok(()) +} + +/// Get the parent commit SHA (HEAD~1) from a worktree. +pub async fn get_parent_sha(worktree_path: &Path) -> Result<String, PatchError> { + let output = Command::new("git") + .current_dir(worktree_path) + .args(["rev-parse", "HEAD~1"]) + .output() + .await + .map_err(|e| PatchError::GitCommand(format!("Failed to get parent SHA: {}", e)))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(PatchError::GitCommand(format!( + "git rev-parse HEAD~1 failed: {}", + stderr + ))); + } + + Ok(String::from_utf8_lossy(&output.stdout).trim().to_string()) +} + +/// Checkout a specific commit in the worktree. +pub async fn checkout_commit(worktree_path: &Path, sha: &str) -> Result<(), PatchError> { + let output = Command::new("git") + .current_dir(worktree_path) + .args(["checkout", sha]) + .output() + .await + .map_err(|e| PatchError::GitCommand(format!("Failed to checkout: {}", e)))?; + + if !output.status.success() { + let stderr = String::from_utf8_lossy(&output.stderr); + return Err(PatchError::GitCommand(format!( + "git checkout {} failed: {}", + sha, stderr + ))); + } + + Ok(()) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::fs; + use tempfile::TempDir; + + async fn setup_test_repo() -> TempDir { + let dir = TempDir::new().unwrap(); + let path = dir.path(); + + // Initialize git repo + Command::new("git") + .current_dir(path) + .args(["init"]) + .output() + .await + .unwrap(); + + // Configure git user + Command::new("git") + .current_dir(path) + .args(["config", "user.email", "test@test.com"]) + .output() + .await + .unwrap(); + Command::new("git") + .current_dir(path) + .args(["config", "user.name", "Test"]) + .output() + .await + .unwrap(); + + // Create initial commit + fs::write(path.join("file.txt"), "initial").unwrap(); + Command::new("git") + .current_dir(path) + .args(["add", "."]) + .output() + .await + .unwrap(); + Command::new("git") + .current_dir(path) + .args(["commit", "-m", "initial"]) + .output() + .await + .unwrap(); + + dir + } + + #[tokio::test] + async fn test_create_and_apply_patch() { + let dir = setup_test_repo().await; + let path = dir.path(); + + // Get base SHA + let base_sha = get_parent_sha(path).await; + // This will fail since there's only one commit + assert!(base_sha.is_err()); + + // Make another commit first + fs::write(path.join("file.txt"), "modified").unwrap(); + Command::new("git") + .current_dir(path) + .args(["add", "."]) + .output() + .await + .unwrap(); + Command::new("git") + .current_dir(path) + .args(["commit", "-m", "modified"]) + .output() + .await + .unwrap(); + + // Now get the base SHA + let base_sha = get_parent_sha(path).await.unwrap(); + + // Create patch + let (patch_data, files_count) = create_patch(path, &base_sha).await.unwrap(); + assert!(!patch_data.is_empty()); + assert_eq!(files_count, 1); + + // Reset to base and apply patch + checkout_commit(path, &base_sha).await.unwrap(); + assert_eq!(fs::read_to_string(path.join("file.txt")).unwrap(), "initial"); + + apply_patch(path, &patch_data).await.unwrap(); + assert_eq!( + fs::read_to_string(path.join("file.txt")).unwrap(), + "modified" + ); + } + + #[tokio::test] + async fn test_empty_patch() { + let dir = setup_test_repo().await; + let path = dir.path(); + + // Make another commit + fs::write(path.join("file.txt"), "modified").unwrap(); + Command::new("git") + .current_dir(path) + .args(["add", "."]) + .output() + .await + .unwrap(); + Command::new("git") + .current_dir(path) + .args(["commit", "-m", "modified"]) + .output() + .await + .unwrap(); + + // Get current HEAD + let head_output = Command::new("git") + .current_dir(path) + .args(["rev-parse", "HEAD"]) + .output() + .await + .unwrap(); + let head_sha = String::from_utf8_lossy(&head_output.stdout) + .trim() + .to_string(); + + // Try to create patch from HEAD to HEAD (no changes) + let result = create_patch(path, &head_sha).await; + assert!(matches!(result, Err(PatchError::EmptyPatch))); + } +} |
