Agent: Deduplicate file uploads using sha1 digest (#1264)

Author: coeuvre

agent/Cargo.lock

@@ -12,6 +12,7 @@ clap = "2.33"
 structopt = "0.3"
 rand = "0.8"
 url = "2.2"
+sha-1 = "0.9"
 
 [dev-dependencies]
 assert_cmd = "2.0"

agent/Cargo.toml

@@ -1,9 +1,11 @@
 use anyhow::{anyhow, Context, Result};
 use serde_json::Value;
+use sha1::{Digest, Sha1};
 use std::{
+    collections::HashSet,
     env,
     fs::{self, File},
-    io::{BufRead, BufReader, Seek, SeekFrom},
+    io::{BufRead, BufReader, Read, Seek, SeekFrom},
     path::{Path, PathBuf, MAIN_SEPARATOR},
     process,
     thread::sleep,
@@ -60,6 +62,7 @@ fn watch_bep_json_file(
     let max_retries = 5;
     let mut retries = max_retries;
     let mut last_offset = 0;
+    let mut uploader = Uploader::new();
 
     'parse_loop: loop {
         match parser.parse() {
@@ -76,9 +79,13 @@ fn watch_bep_json_file(
                         .filter(|test_result| status.contains(&test_result.overall_status.as_str()))
                     {
                         for failed_test in test_summary.failed.iter() {
-                            if let Err(error) =
-                                upload_test_log(dry, local_exec_root, &failed_test.uri, mode)
-                            {
+                            if let Err(error) = upload_test_log(
+                                &mut uploader,
+                                dry,
+                                local_exec_root,
+                                &failed_test.uri,
+                                mode,
+                            ) {
                                 error!("{:?}", error);
                             }
                         }
@@ -107,16 +114,21 @@ fn watch_bep_json_file(
     let should_upload_bep_json_file =
         debug || (monitor_flaky_tests && parser.has_overall_test_status("FLAKY"));
     if should_upload_bep_json_file {
-        if let Err(error) = upload_bep_json_file(dry, build_event_json_file, mode) {
+        if let Err(error) = upload_bep_json_file(&mut uploader, dry, build_event_json_file, mode) {
             error!("{:?}", error);
         }
     }
 
     Ok(())
 }
 
-fn upload_bep_json_file(dry: bool, build_event_json_file: &Path, mode: Mode) -> Result<()> {
-    upload_artifact(dry, None, build_event_json_file, mode)
+fn upload_bep_json_file(
+    uploader: &mut Uploader,
+    dry: bool,
+    build_event_json_file: &Path,
+    mode: Mode,
+) -> Result<()> {
+    uploader.upload_artifact(dry, None, build_event_json_file, mode)
 }
 
 fn execute_command(dry: bool, cwd: Option<&Path>, program: &str, args: &[&str]) -> Result<()> {
@@ -148,19 +160,75 @@ fn execute_command(dry: bool, cwd: Option<&Path>, program: &str, args: &[&str])
     Ok(())
 }
 
-fn upload_artifact_buildkite(dry: bool, cwd: Option<&Path>, artifact: &Path) -> Result<()> {
-    let artifact = artifact.display().to_string();
-    execute_command(
-        dry,
-        cwd,
-        "buildkite-agent",
-        &["artifact", "upload", artifact.as_str()],
-    )
+type Sha1Digest = [u8; 20];
+
+fn read_entire_file(path: &Path) -> Result<Vec<u8>> {
+    let mut file = File::open(path)?;
+    let mut buf = Vec::new();
+    file.read_to_end(&mut buf)?;
+    Ok(buf)
+}
+
+fn sha1_digest(path: &Path) -> Sha1Digest {
+    let buf = match read_entire_file(path) {
+        Ok(buf) => buf,
+        _ => path.display().to_string().into_bytes(),
+    };
+
+    let mut hasher = Sha1::new();
+    hasher.update(buf);
+    let hash = hasher.finalize();
+    hash.into()
+}
+
+struct Uploader {
+    uploaded_digests: HashSet<Sha1Digest>,
 }
 
-fn upload_artifact(dry: bool, cwd: Option<&Path>, artifact: &Path, mode: Mode) -> Result<()> {
-    match mode {
-        Mode::Buildkite => upload_artifact_buildkite(dry, cwd, artifact),
+impl Uploader {
+    pub fn new() -> Self {
+        Self {
+            uploaded_digests: HashSet::new(),
+        }
+    }
+
+    pub fn upload_artifact(
+        &mut self,
+        dry: bool,
+        cwd: Option<&Path>,
+        artifact: &Path,
+        mode: Mode,
+    ) -> Result<()> {
+        {
+            let file = match cwd {
+                Some(cwd) => cwd.join(artifact),
+                None => PathBuf::from(artifact),
+            };
+            let digest = sha1_digest(&file);
+            if self.uploaded_digests.contains(&digest) {
+                return Ok(());
+            }
+            self.uploaded_digests.insert(digest);
+        }
+
+        match mode {
+            Mode::Buildkite => self.upload_artifact_buildkite(dry, cwd, artifact),
+        }
+    }
+
+    fn upload_artifact_buildkite(
+        &mut self,
+        dry: bool,
+        cwd: Option<&Path>,
+        artifact: &Path,
+    ) -> Result<()> {
+        let artifact = artifact.display().to_string();
+        execute_command(
+            dry,
+            cwd,
+            "buildkite-agent",
+            &["artifact", "upload", artifact.as_str()],
+        )
     }
 }
 
@@ -213,6 +281,7 @@ fn uri_to_file_path(uri: &str) -> Result<PathBuf> {
 }
 
 fn upload_test_log(
+    uploader: &mut Uploader,
     dry: bool,
     local_exec_root: Option<&Path>,
     test_log: &str,
@@ -221,7 +290,7 @@ fn upload_test_log(
     let path = uri_to_file_path(test_log)?;
 
     if let Some((first, second)) = split_path_inclusive(&path, "testlogs") {
-        return upload_artifact(dry, Some(&first), &second, mode);
+        return uploader.upload_artifact(dry, Some(&first), &second, mode);
     }
 
     let artifact = if let Some(local_exec_root) = local_exec_root {
@@ -234,7 +303,7 @@ fn upload_test_log(
         &path
     };
 
-    upload_artifact(dry, local_exec_root, &artifact, mode)
+    uploader.upload_artifact(dry, local_exec_root, &artifact, mode)
 }
 
 #[derive(Debug)]

agent/src/artifact/upload.rs

@@ -125,3 +125,21 @@ fn truncate_build_event_json_file_recover_from_middle() -> Result<()> {
         Ok(())
     })
 }
+
+#[cfg(not(target_os = "windows"))]
+#[test]
+fn test_logs_deduplicated() -> Result<()> {
+    let mut cmd = Command::cargo_bin("bazelci-agent")?;
+    cmd.args([
+        "artifact",
+        "upload",
+        "--dry",
+        "--mode=buildkite",
+        "--build_event_json_file=tests/data/test_bep_duplicated.json",
+    ]);
+    cmd.assert()
+        .success()
+        .stdout(predicates::str::contains("buildkite-agent artifact upload src/test/shell/bazel/starlark_repository_test/shard_4_of_6/test_attempts/attempt_1.log").count(1));
+
+    Ok(())
+}

agent/tests/artifact/upload.rs

@@ -0,0 +1,4 @@
+{"id":{"workspace":{}},"workspaceInfo":{"localExecRoot":"/private/var/tmp/_bazel_buildkite/78a0792bf9bb0133b1a4a7d083181fcb/execroot/io_bazel"}}
+{"id":{"testSummary":{"label":"//src/test/shell/bazel:starlark_repository_test","configuration":{"id":"7479eaa1eeb472e5c3fdd9f0b604289ffbe45a36edb8a7f474df0c95501b4d00"}}},"testSummary":{"totalRunCount":7,"failed":[{"uri":"file:///private/var/tmp/_bazel_buildkite/78a0792bf9bb0133b1a4a7d083181fcb/execroot/io_bazel/bazel-out/darwin-fastbuild/testlogs/src/test/shell/bazel/starlark_repository_test/shard_4_of_6/test_attempts/attempt_1.log"}],"overallStatus":"FLAKY","firstStartTimeMillis":"1630444947193","lastStopTimeMillis":"1630445154997","totalRunDurationMillis":"338280","runCount":1,"shardCount":6}}
+{"id":{"testSummary":{"label":"//src/test/shell/bazel:starlark_repository_test","configuration":{"id":"7479eaa1eeb472e5c3fdd9f0b604289ffbe45a36edb8a7f474df0c95501b4d00"}}},"testSummary":{"totalRunCount":7,"failed":[{"uri":"file:///private/var/tmp/_bazel_buildkite/78a0792bf9bb0133b1a4a7d083181fcb/execroot/io_bazel/bazel-out/darwin-fastbuild/testlogs/src/test/shell/bazel/starlark_repository_test/shard_4_of_6/test_attempts/attempt_1.log"}],"overallStatus":"FLAKY","firstStartTimeMillis":"1630444947193","lastStopTimeMillis":"1630445154997","totalRunDurationMillis":"338280","runCount":1,"shardCount":6}}
+{"id":{"progress":{}},"progress":{},"lastMessage":true}