allow for more coverage data in total (#519)

src/agent/onefuzz-agent/src/tasks/coverage/libfuzzer_coverage.rs

@@ -43,7 +43,7 @@ use onefuzz::{fs::list_files, libfuzzer::LibFuzzer, syncdir::SyncedDir};
 use onefuzz_telemetry::{Event::coverage_data, EventData};
 use reqwest::Url;
 use serde::Deserialize;
-use std::collections::HashMap;
+use std::collections::{BTreeMap, HashMap};
 use std::{
     ffi::OsString,
     path::{Path, PathBuf},
@@ -197,7 +197,7 @@ pub struct CoverageProcessor {
     config: Arc<Config>,
     pub recorder: CoverageRecorder,
     pub total: TotalCoverage,
-    pub module_totals: HashMap<OsString, TotalCoverage>,
+    pub module_totals: BTreeMap<OsString, TotalCoverage>,
     heartbeat_client: Option<TaskHeartbeatClient>,
 }
 
@@ -206,7 +206,7 @@ impl CoverageProcessor {
         let heartbeat_client = config.common.init_heartbeat().await?;
         let total = TotalCoverage::new(config.coverage.path.join(TOTAL_COVERAGE));
         let recorder = CoverageRecorder::new(config.clone());
-        let module_totals = HashMap::default();
+        let module_totals = BTreeMap::default();
 
         Ok(Self {
             config,
@@ -244,8 +244,9 @@ impl CoverageProcessor {
         Ok(())
     }
 
-    async fn collect_by_module(&mut self, path: &Path) -> Result<PathBuf> {
-        let files = list_files(&path).await?;
+    async fn collect_by_module(&mut self, path: &Path) -> Result<()> {
+        let mut files = list_files(&path).await?;
+        files.sort();
         let mut sum = Vec::new();
 
         for file in &files {
@@ -264,14 +265,25 @@ impl CoverageProcessor {
             .await
             .with_context(|| format!("unable to write combined coverage file: {:?}", combined))?;
 
-        Ok(combined.into())
+        Ok(())
     }
 
     pub async fn test_input(&mut self, input: &Path) -> Result<()> {
         info!("processing input {:?}", input);
         let new_coverage = self.recorder.record(input).await?;
-        let combined = self.collect_by_module(&new_coverage).await?;
-        self.total.update(&combined).await?;
+        self.collect_by_module(&new_coverage).await?;
+        self.update_total().await?;
+        Ok(())
+    }
+
+    async fn update_total(&mut self) -> Result<()> {
+        let mut total = Vec::new();
+        for module_total in self.module_totals.values() {
+            if let Some(mut module_data) = module_total.data().await? {
+                total.append(&mut module_data);
+            }
+        }
+        self.total.write(&total).await?;
         Ok(())
     }
 

src/agent/onefuzz-agent/src/tasks/coverage/total.rs

@@ -43,22 +43,29 @@ impl TotalCoverage {
         &self.path
     }
 
+    pub async fn write(&self, data: &[u8]) -> Result<()> {
+        fs::write(self.path(), data).await?;
+        Ok(())
+    }
+
     pub async fn update_bytes(&self, new_data: &[u8]) -> Result<()> {
         match self.data().await {
             Ok(Some(mut total_data)) => {
+                if total_data.len() < new_data.len() {
+                    total_data.resize_with(new_data.len(), || 0);
+                }
                 for (i, b) in new_data.iter().enumerate() {
                     if *b > 0 {
                         total_data[i] = 1;
                     }
                 }
-
-                fs::write(self.path(), total_data).await?;
+                self.write(&total_data).await?;
             }
             Ok(None) => {
                 // Base case: we don't yet have any total coverage. Promote the
                 // new coverage to being our total coverage.
                 info!("initializing total coverage map {}", self.path().display());
-                fs::write(self.path(), new_data).await?;
+                self.write(new_data).await?;
             }
             Err(err) => {
                 // Couldn't read total for some other reason, so this is a real error.
@@ -69,11 +76,6 @@ impl TotalCoverage {
         Ok(())
     }
 
-    pub async fn update(&self, new: impl AsRef<Path>) -> Result<()> {
-        let new_data = fs::read(new).await?;
-        self.update_bytes(&new_data).await
-    }
-
     pub async fn info(&self) -> Result<Info> {
         let data = self
             .data()