Incrementally make progress during collection

This permits restarting of the collection server without losing hour(s) of work.
It also removes the extremely long-lived (2h30m as of now) transaction in favor
of one wrapping each individual benchmark. Ideally, that would be by itself
"sufficiently small."

However, currently, even just limiting to one benchmark can mean a hour-long
transaction (for script-servo-2). This is a bit unfortunate, but at least
killing that transaction will be dealt with fairly smoothly.

Author: Mark-Simulacrum

collector/src/main.rs

@@ -195,14 +195,6 @@ fn bench(
     self_profile: bool,
 ) -> BenchmarkErrors {
     let mut conn = rt.block_on(pool.connection());
-    let status_conn;
-    let status_conn: Option<&dyn database::Connection> =
-        if conn.separate_transaction_for_collector() {
-            status_conn = rt.block_on(pool.connection());
-            Some(&*status_conn)
-        } else {
-            None
-        };
     let mut errors = BenchmarkErrors::new();
     eprintln!("Benchmarking {} for triple {}", cid, compiler.triple);
 
@@ -215,38 +207,24 @@ fn bench(
         );
     }
 
-    let mut tx = rt.block_on(conn.transaction());
-    let index = rt.block_on(database::Index::load(tx.conn()));
-    let has_collected = match cid {
-        ArtifactId::Commit(commit) => index.commits().iter().any(|c| c == commit),
-        ArtifactId::Artifact(id) => index.artifacts().any(|a| a == id),
-    };
-    if has_collected {
-        eprintln!("'{}' has previously been collected, aborting.", cid);
-        eprintln!(
-            "Note that this behavior is likely to change in the future \
-            to collect and append the data instead."
-        );
-        return errors;
-    }
-    let interned_cid = rt.block_on(tx.conn().artifact_id(&cid));
+    let interned_cid = rt.block_on(conn.artifact_id(&cid));
 
     let start = Instant::now();
     let steps = benchmarks
         .iter()
         .map(|b| b.name.to_string())
         .collect::<Vec<_>>();
-    rt.block_on(
-        status_conn
-            .unwrap_or_else(|| &*tx.conn())
-            .collector_start(interned_cid, &steps),
-    );
+    rt.block_on(conn.collector_start(interned_cid, &steps));
+    let mut skipped = false;
     for (nth_benchmark, benchmark) in benchmarks.iter().enumerate() {
-        rt.block_on(
-            status_conn
-                .unwrap_or_else(|| &*tx.conn())
-                .collector_start_step(interned_cid, &benchmark.name.to_string()),
-        );
+        let is_fresh =
+            rt.block_on(conn.collector_start_step(interned_cid, &benchmark.name.to_string()));
+        if !is_fresh {
+            skipped = true;
+            eprintln!("skipping {} -- already benchmarked", benchmark.name);
+            continue;
+        }
+        let mut tx = rt.block_on(conn.transaction());
         rt.block_on(
             tx.conn()
                 .record_benchmark(benchmark.name.0.as_str(), benchmark.supports_stable()),
@@ -278,19 +256,19 @@ fn bench(
             ));
         };
         rt.block_on(
-            status_conn
-                .unwrap_or_else(|| &*tx.conn())
+            tx.conn()
                 .collector_end_step(interned_cid, &benchmark.name.to_string()),
         );
+        rt.block_on(tx.commit()).expect("committed");
     }
     let end = start.elapsed();
 
     eprintln!("collection took {:?}", end);
 
-    rt.block_on(tx.conn().record_duration(interned_cid, end));
-
-    // Publish results now that we've finished fully with this commit.
-    rt.block_on(tx.commit()).unwrap();
+    if !skipped {
+        log::info!("skipping duration record -- skipped parts of run");
+        rt.block_on(conn.record_duration(interned_cid, end));
+    }
 
     rt.block_on(async move {
         // This ensures that we're good to go with the just updated data.

database/src/pool.rs

@@ -65,20 +65,11 @@ pub trait Connection: Send + Sync {
     // Collector status API
 
     async fn collector_start(&self, aid: ArtifactIdNumber, steps: &[String]);
-    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str);
-    async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str);
 
-    // This returns `true` if the collector commands can be placed in a separate
-    // transaction.
-    //
-    // Currently, the sqlite backend does not support "regular" usage where they
-    // are used for genuine progress reporting. sqlite does not support
-    // concurrent writers -- it will return an error (or wait, if a busy timeout
-    // is configured).
-    //
-    // For now we don't care much as sqlite is not used in production and in
-    // local usage you can just look at the logs.
-    fn separate_transaction_for_collector(&self) -> bool;
+    // Returns `true` if the step was started, i.e., it did not previously have
+    // an end. Otherwise returns false, indicating that we can skip it.
+    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) -> bool;
+    async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str);
 }
 
 #[async_trait::async_trait]

database/src/pool/postgres.rs

@@ -167,6 +167,7 @@ static MIGRATIONS: &[&str] = &[
         end_time timestamptz
     );
     "#,
+    r#"alter table collector_progress add unique (aid, step);"#,
 ];
 
 #[async_trait::async_trait]
@@ -781,6 +782,7 @@ where
     }
 
     async fn collector_start(&self, aid: ArtifactIdNumber, steps: &[String]) {
+        // Clean up -- we'll re-insert any missing things in the loop below.
         self.conn()
             .execute(
                 "delete from collector_progress where start_time is null or end_time is null;",
@@ -792,34 +794,40 @@ where
         for step in steps {
             self.conn()
                 .execute(
-                    "insert into collector_progress(aid, step) VALUES ($1, $2)",
+                    "insert into collector_progress(aid, step) VALUES ($1, $2)
+                    ON CONFLICT DO NOTHING",
                     &[&(aid.0 as i16), &step],
                 )
                 .await
                 .unwrap();
         }
     }
-    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) {
+    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) -> bool {
+        // If we modified a row, then we populated a start time, so we're good
+        // to go. Otherwise we should just skip this step.
         self.conn()
             .execute(
                 "update collector_progress set start_time = statement_timestamp() \
-                where aid = $1 and step = $2 and start_time is null and end_time is null;",
+                where aid = $1 and step = $2 and end_time is null;",
                 &[&(aid.0 as i16), &step],
             )
             .await
-            .unwrap();
+            .unwrap()
+            == 1
     }
     async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str) {
-        self.conn()
+        let did_modify = self
+            .conn()
             .execute(
                 "update collector_progress set end_time = statement_timestamp() \
                 where aid = $1 and step = $2 and start_time is not null and end_time is null;",
                 &[&(aid.0 as i16), &step],
             )
             .await
-            .unwrap();
-    }
-    fn separate_transaction_for_collector(&self) -> bool {
-        true
+            .unwrap()
+            == 1;
+        if !did_modify {
+            log::error!("did not end {} for {:?}", step, aid);
+        }
     }
 }

database/src/pool/sqlite.rs

@@ -145,7 +145,8 @@ static MIGRATIONS: &[&str] = &[
         aid integer not null references artifact(id) on delete cascade on update cascade,
         step text not null,
         start integer,
-        end integer
+        end integer,
+        UNIQUE(aid, step)
     );
     "#,
 ];
@@ -227,7 +228,7 @@ impl Connection for SqliteConnection {
     async fn record_duration(&self, artifact: ArtifactIdNumber, duration: Duration) {
         self.raw_ref()
             .prepare_cached(
-                "insert into artifact_collection_duration (aid, date_recorded, duration) VALUES (?, strftime('%s','now'), ?)",
+                "insert or ignore into artifact_collection_duration (aid, date_recorded, duration) VALUES (?, strftime('%s','now'), ?)",
             )
             .unwrap()
             .execute(params![artifact.0, duration.as_secs() as i64])
@@ -644,31 +645,34 @@ impl Connection for SqliteConnection {
         for step in steps {
             self.raw_ref()
                 .execute(
-                    "insert into collector_progress(aid, step) VALUES (?, ?)",
+                    "insert or ignore into collector_progress(aid, step) VALUES (?, ?)",
                     params![&aid.0, step],
                 )
                 .unwrap();
         }
     }
-    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) {
+    async fn collector_start_step(&self, aid: ArtifactIdNumber, step: &str) -> bool {
         self.raw_ref()
             .execute(
                 "update collector_progress set start = strftime('%s','now') \
-            where aid = ? and step = ? and start is null;",
+                where aid = ? and step = ? and end is null;",
                 params![&aid.0, &step],
             )
-            .unwrap();
+            .unwrap()
+            == 1
     }
     async fn collector_end_step(&self, aid: ArtifactIdNumber, step: &str) {
-        self.raw_ref()
+        let did_modify = self
+            .raw_ref()
             .execute(
                 "update collector_progress set end = strftime('%s','now') \
-            where aid = ? and step = ? and start is not null and end is null;",
+                where aid = ? and step = ? and start is not null and end is null;",
                 params![&aid.0, &step],
             )
-            .unwrap();
-    }
-    fn separate_transaction_for_collector(&self) -> bool {
-        false
+            .unwrap()
+            == 1;
+        if !did_modify {
+            log::error!("did not end {} for {:?}", step, aid);
+        }
     }
 }