Rework getting next crate to be per-worker, rather than a global queue

The previous commit added the notion of timeouts on retrieved crates from the
server, which meant that when an agent (which has multiple workers) requested
the 1024 crates from the server, the tail of that request would likely have
timed out by the time the agent started/completed working on those crates. This
isn't technically a problem, but it is likely to lead to duplicated work for
that tail, so we avoid it by retrieving crates one at a time with the new
commit.

Author: Mark-Simulacrum

src/agent/api.rs

@@ -120,15 +120,15 @@ impl AgentApi {
         })
     }
 
-    pub fn next_experiment(&self) -> Fallible<(Experiment, Vec<Crate>)> {
+    pub fn next_experiment(&self) -> Fallible<Experiment> {
         self.retry(|this| loop {
             let resp: Option<_> = this
-                .build_request(Method::GET, "next-experiment")
+                .build_request(Method::POST, "next-experiment")
                 .send()?
                 .to_api_response()?;
 
-            if let Some((experiment, crates)) = resp {
-                return Ok((experiment, crates));
+            if let Some(experiment) = resp {
+                return Ok(experiment);
             }
 
             // If we're just waiting for an experiment, we should be considered
@@ -139,6 +139,15 @@ impl AgentApi {
         })
     }
 
+    pub fn next_crate(&self, ex: &str) -> Fallible<Option<Crate>> {
+        self.retry(|this| loop {
+            this.build_request(Method::POST, "next-crate")
+                .json(&json!(ex))
+                .send()?
+                .to_api_response()?
+        })
+    }
+
     pub fn record_progress(
         &self,
         ex: &Experiment,

src/agent/mod.rs

@@ -69,9 +69,9 @@ impl FromIterator<String> for Capabilities {
     }
 }
 
-struct Agent {
+pub struct Agent {
     api: AgentApi,
-    config: Config,
+    pub config: Config,
 }
 
 impl Agent {
@@ -90,10 +90,14 @@ impl Agent {
         })
     }
 
-    fn experiment(&self) -> Fallible<(Experiment, Vec<Crate>)> {
+    fn experiment(&self) -> Fallible<Experiment> {
         info!("asking the server for a new experiment...");
         Ok(self.api.next_experiment()?)
     }
+
+    pub fn next_crate(&self, ex: &str) -> Fallible<Option<Crate>> {
+        self.api.next_crate(ex)
+    }
 }
 
 static HEALTH_CHECK: AtomicBool = AtomicBool::new(false);
@@ -150,7 +154,7 @@ fn run_experiment(
     threads_count: usize,
     past_experiment: &mut Option<String>,
 ) -> Result<(), (Option<Experiment>, Error)> {
-    let (ex, crates) = agent.experiment().map_err(|e| (None, e))?;
+    let ex = agent.experiment().map_err(|e| (None, e))?;
 
     if Some(&ex.name) != past_experiment.as_ref() {
         debug!("purging build directories...");
@@ -170,8 +174,10 @@ fn run_experiment(
         }
     }
 
-    crate::runner::run_ex(&ex, workspace, &crates, db, threads_count, &agent.config)
-        .map_err(|err| (Some(ex), err))?;
+    crate::runner::run_ex(&ex, workspace, db, threads_count, &agent.config, &|| {
+        agent.next_crate(&ex.name)
+    })
+    .map_err(|err| (Some(ex), err))?;
     Ok(())
 }
 

src/cli.rs

@@ -468,13 +468,16 @@ impl Crater {
                     let workspace = self
                         .workspace(docker_env.as_ref().map(|s| s.as_str()), fast_workspace_init)?;
                     workspace.purge_all_build_dirs()?;
+
+                    let crates =
+                        std::sync::Mutex::new(experiment.get_uncompleted_crates(&db, None)?);
                     let res = runner::run_ex(
                         &experiment,
                         &workspace,
-                        &experiment.get_uncompleted_crates(&db, &config)?,
                         &result_db,
                         threads,
                         &config,
+                        &|| Ok(crates.lock().unwrap().pop()),
                     );
                     workspace.purge_all_build_dirs()?;
                     res?;

src/db/mod.rs

@@ -58,6 +58,7 @@ impl Database {
     #[cfg(test)]
     pub fn temp() -> Fallible<Self> {
         let tempfile = NamedTempFile::new()?;
+        dbg!(&tempfile.path());
         Database::new(
             SqliteConnectionManager::file(tempfile.path()),
             Some(tempfile),

src/experiments.rs

@@ -1,4 +1,3 @@
-use crate::config::Config;
 use crate::crates::Crate;
 use crate::db::{Database, QueryUtils};
 use crate::prelude::*;
@@ -13,7 +12,6 @@ use std::str::FromStr;
 use url::Url;
 
 //sqlite limit is ignored if the expression evaluates to a negative value
-static FULL_LIST: i32 = -1;
 static SQL_VARIABLE_LIMIT: usize = 500;
 
 string_enum!(pub enum Status {
@@ -595,18 +593,12 @@ impl Experiment {
         .collect::<Fallible<Vec<Crate>>>()
     }
 
-    fn crate_list_size(&self, config: &Config) -> i32 {
-        match self.assigned_to {
-            //if experiment is distributed return chunk
-            Some(Assignee::Distributed) => config.chunk_size(),
-            //if experiment is assigned to specific agent return all the crates
-            _ => FULL_LIST,
-        }
-    }
-
-    pub fn get_uncompleted_crates(&self, db: &Database, config: &Config) -> Fallible<Vec<Crate>> {
-        let limit = self.crate_list_size(config);
-
+    pub fn get_uncompleted_crates(
+        &self,
+        db: &Database,
+        limit: Option<u32>,
+    ) -> Fallible<Vec<Crate>> {
+        let limit = limit.map(|l| l as i32).unwrap_or(-1);
         #[cfg(not(test))]
         const RUN_TIMEOUT: u32 = 20;
         #[cfg(test)]
@@ -999,12 +991,12 @@ mod tests {
         // Create a dummy experiment
         CreateExperiment::dummy("dummy").apply(&ctx).unwrap();
         let ex = Experiment::get(&db, "dummy").unwrap().unwrap();
-        let crates = ex.get_uncompleted_crates(&db, &config).unwrap();
+        let crates = ex.get_uncompleted_crates(&db, None).unwrap();
         // Assert the whole list is returned
         assert_eq!(crates.len(), ex.get_crates(&db).unwrap().len());
 
         // Test already completed crates does not show up again
-        let uncompleted_crates = ex.get_uncompleted_crates(&db, &config).unwrap();
+        let uncompleted_crates = ex.get_uncompleted_crates(&db, None).unwrap();
         assert_eq!(uncompleted_crates.len(), 0);
     }
 
@@ -1022,10 +1014,10 @@ mod tests {
         // Create a dummy experiment
         CreateExperiment::dummy("dummy").apply(&ctx).unwrap();
         let ex = Experiment::next(&db, &agent1).unwrap().unwrap().1;
-        assert!(!ex.get_uncompleted_crates(&db, &config).unwrap().is_empty());
+        assert!(!ex.get_uncompleted_crates(&db, None).unwrap().is_empty());
         assert!(Experiment::next(&db, &agent1).unwrap().is_some());
         std::thread::sleep(std::time::Duration::from_secs(80)); // need to wait for at least 60 seconds for timeout to fire
         assert_eq!(ex.status, Status::Running);
-        assert!(!ex.get_uncompleted_crates(&db, &config).unwrap().is_empty());
+        assert!(!ex.get_uncompleted_crates(&db, None).unwrap().is_empty());
     }
 }

src/runner/mod.rs

@@ -48,10 +48,10 @@ impl RunnerState {
 pub fn run_ex<DB: WriteResults + Sync>(
     ex: &Experiment,
     workspace: &Workspace,
-    crates: &[Crate],
     db: &DB,
     threads_count: usize,
     config: &Config,
+    next_crate: &(dyn Fn() -> Fallible<Option<Crate>> + Send + Sync),
 ) -> Fallible<()> {
     // Attempt to spin indefinitely until docker is up. Ideally, we would
     // decomission this agent until docker is up, instead of leaving the
@@ -105,17 +105,16 @@ pub fn run_ex<DB: WriteResults + Sync>(
 
     let state = RunnerState::new();
 
-    let crates = Mutex::new(crates.to_vec());
     let workers = (0..threads_count)
         .map(|i| {
             Worker::new(
                 format!("worker-{}", i),
                 workspace,
                 ex,
                 config,
-                &crates,
                 &state,
                 db,
+                next_crate,
             )
         })
         .collect::<Vec<_>>();

src/runner/worker.rs

@@ -1,4 +1,3 @@
-use crate::config::Config;
 use crate::crates::Crate;
 use crate::experiments::{Experiment, Mode};
 use crate::prelude::*;
@@ -19,30 +18,30 @@ pub(super) struct Worker<'a, DB: WriteResults + Sync> {
     workspace: &'a Workspace,
     build_dir: Mutex<BuildDirectory>,
     ex: &'a Experiment,
-    config: &'a Config,
-    crates: &'a Mutex<Vec<Crate>>,
+    config: &'a crate::config::Config,
     state: &'a RunnerState,
     db: &'a DB,
     target_dir_cleanup: AtomicBool,
+    next_crate: &'a (dyn Fn() -> Fallible<Option<Crate>> + Send + Sync),
 }
 
 impl<'a, DB: WriteResults + Sync> Worker<'a, DB> {
     pub(super) fn new(
         name: String,
         workspace: &'a Workspace,
         ex: &'a Experiment,
-        config: &'a Config,
-        crates: &'a Mutex<Vec<Crate>>,
+        config: &'a crate::config::Config,
         state: &'a RunnerState,
         db: &'a DB,
+        next_crate: &'a (dyn Fn() -> Fallible<Option<Crate>> + Send + Sync),
     ) -> Self {
         Worker {
             build_dir: Mutex::new(workspace.build_dir(&name)),
             name,
             workspace,
             ex,
             config,
-            crates,
+            next_crate,
             state,
             db,
             target_dir_cleanup: AtomicBool::new(false),
@@ -92,7 +91,7 @@ impl<'a, DB: WriteResults + Sync> Worker<'a, DB> {
 
     pub(super) fn run(&self) -> Fallible<()> {
         loop {
-            let krate = if let Some(next) = self.crates.lock().unwrap().pop() {
+            let krate = if let Some(next) = (self.next_crate)()? {
                 next
             } else {
                 // We're done if no more crates left.

src/server/agents.rs

@@ -28,6 +28,11 @@ pub struct Agent {
 impl Agent {
     fn with_experiment(mut self, db: &Database) -> Fallible<Self> {
         self.experiment = Experiment::run_by(db, &Assignee::Agent(self.name.clone()))?;
+        eprintln!(
+            "{} has experiment {:?}",
+            self.name,
+            self.experiment.as_ref().map(|e| &e.name)
+        );
         Ok(self)
     }
 
@@ -278,10 +283,12 @@ mod tests {
         let (_new, ex) = Experiment::next(&db, &Assignee::Agent("agent".to_string()))
             .unwrap()
             .unwrap();
-        ex.get_uncompleted_crates(&db, &config).unwrap();
+        ex.get_uncompleted_crates(&db, None).unwrap();
 
         // After an experiment is assigned to the agent, the agent is working
         let agent = agents.get("agent").unwrap().unwrap();
+        dbg!(agent.status());
+        std::thread::sleep(std::time::Duration::from_secs(100000));
         assert_eq!(agent.status(), AgentStatus::Working);
     }
 

src/server/metrics.rs

@@ -245,7 +245,7 @@ mod tests {
         .apply(&ctx)
         .unwrap();
         let ex = Experiment::next(&db, &assignee).unwrap().unwrap().1;
-        ex.get_uncompleted_crates(&db, &config).unwrap();
+        ex.get_uncompleted_crates(&db, None).unwrap();
         METRICS.update_agent_status(&db, &agent_list_ref).unwrap();
 
         // There are no experiments in the queue but agent1 is still executing the

src/server/routes/agent.rs

@@ -40,14 +40,22 @@ pub fn routes(
         .and(auth_filter(data.clone(), TokenType::Agent))
         .map(endpoint_config);
 
-    let next_experiment = warp::get2()
+    let next_experiment = warp::post2()
         .and(warp::path("next-experiment"))
         .and(warp::path::end())
         .and(mutex_filter.clone())
         .and(github_data_filter)
         .and(auth_filter(data.clone(), TokenType::Agent))
         .map(endpoint_next_experiment);
 
+    let next_crate = warp::post2()
+        .and(warp::path("next-crate"))
+        .and(warp::path::end())
+        .and(warp::body::json())
+        .and(data_filter.clone())
+        .and(auth_filter(data.clone(), TokenType::Agent))
+        .map(endpoint_next_crate);
+
     let record_progress = warp::post2()
         .and(warp::path("record-progress"))
         .and(warp::path::end())
@@ -76,6 +84,8 @@ pub fn routes(
             config
                 .or(next_experiment)
                 .unify()
+                .or(next_crate)
+                .unify()
                 .or(record_progress)
                 .unify()
                 .or(heartbeat)
@@ -126,10 +136,26 @@ fn endpoint_next_experiment(
             }
         }
 
-        Some((
-            ex.clone(),
-            ex.get_uncompleted_crates(&data.db, &data.config)?,
-        ))
+        Some(ex)
+    } else {
+        None
+    };
+
+    Ok(ApiResponse::Success { result }.into_response()?)
+}
+
+fn endpoint_next_crate(
+    experiment: String,
+    data: Arc<Data>,
+    _auth: AuthDetails,
+) -> Fallible<Response<Body>> {
+    let result = if let Some(ex) = Experiment::get(&data.db, &experiment)? {
+        let mut crates = ex.get_uncompleted_crates(&data.db, Some(1))?;
+        if crates.is_empty() {
+            None
+        } else {
+            Some(crates.remove(0))
+        }
     } else {
         None
     };