fix(datastore): make legacy Android bucket merge linear instead of quadratic

aw-datastore/src/datastore.rs

@@ -248,6 +248,131 @@ fn parse_event_row(row: &rusqlite::Row, clip: Option<(i64, i64)>) -> rusqlite::R
     })
 }
 
+/// Legacy-bucket events that overlap no other event in either bucket, so they can
+/// be moved into the destination bucket without creating overlapping records.
+///
+/// Two events overlap when `a.starttime < b.endtime AND b.starttime < a.endtime`
+/// (strict, matching the datastore's overlap semantics elsewhere). Events are
+/// read once ordered by `starttime`. A sweep keeps still-open events in a
+/// min-heap keyed by `endtime`. A positive-duration event overlaps every
+/// still-open event, so those are marked in O(1) via an epoch counter rather
+/// than scanning the heap (which would be quadratic for stacked histories).
+/// Zero-duration events are the only case that still scans, and they never stay
+/// in the open set. The pass is O(n log n).
+fn movable_legacy_event_ids(
+    conn: &Connection,
+    legacy_bid: Option<i64>,
+    destination_bid: Option<i64>,
+) -> rusqlite::Result<Vec<i64>> {
+    use std::cmp::Reverse;
+    use std::collections::BinaryHeap;
+
+    struct Span {
+        id: i64,
+        starttime: i64,
+        endtime: i64,
+        legacy: bool,
+    }
+
+    let mut stmt = conn.prepare(
+        "SELECT id, starttime, endtime, bucketrow FROM events
+         WHERE bucketrow IN (?1, ?2)
+         ORDER BY starttime ASC, id ASC",
+    )?;
+    let spans = stmt
+        .query_map(params![legacy_bid, destination_bid], |row| {
+            let bucketrow: Option<i64> = row.get(3)?;
+            Ok(Span {
+                id: row.get(0)?,
+                starttime: row.get(1)?,
+                endtime: row.get(2)?,
+                legacy: bucketrow == legacy_bid,
+            })
+        })?
+        .collect::<rusqlite::Result<Vec<Span>>>()?;
+
+    let mut overlaps = vec![false; spans.len()];
+    // Open events: (endtime, index, epoch-at-push). Epoch marks the whole
+    // concurrent set in O(1) when a later positive-duration event overlaps it.
+    let mut open: BinaryHeap<Reverse<(i64, usize, u32)>> = BinaryHeap::new();
+    let mut epoch: u32 = 0;
+    for (i, span) in spans.iter().enumerate() {
+        // Events that ended at or before this start can never overlap this or
+        // any later event (later events start no earlier than this one).
+        while let Some(Reverse((endtime, j, pushed_epoch))) = open.peek().copied() {
+            if endtime <= span.starttime {
+                open.pop();
+                if pushed_epoch < epoch {
+                    overlaps[j] = true;
+                }
+            } else {
+                break;
+            }
+        }
+        if !open.is_empty() {
+            if span.endtime > span.starttime {
+                // Positive duration: every remaining open event `j` has
+                // `j.starttime <= span.starttime < span.endtime` and
+                // `span.starttime < j.endtime`, so they all overlap.
+                overlaps[i] = true;
+                epoch += 1;
+            } else {
+                // Zero-duration: overlaps open events that started strictly
+                // earlier, but not same-start positive-duration events
+                // (`j.starttime < span.endtime` fails when starttimes match).
+                let mut overlapped = false;
+                for Reverse((_, j, _)) in open.iter() {
+                    if spans[*j].starttime < span.starttime {
+                        overlaps[*j] = true;
+                        overlapped = true;
+                    }
+                }
+                if overlapped {
+                    overlaps[i] = true;
+                }
+            }
+        }
+        open.push(Reverse((span.endtime, i, epoch)));
+    }
+    while let Some(Reverse((_, j, pushed_epoch))) = open.pop() {
+        if pushed_epoch < epoch {
+            overlaps[j] = true;
+        }
+    }
+
+    Ok(spans
+        .iter()
+        .zip(overlaps.iter())
+        .filter(|(span, overlapping)| span.legacy && !**overlapping)
+        .map(|(span, _)| span.id)
+        .collect())
+}
+
+/// Reassign the given events to `bucket_bid`, in batches so the statement stays
+/// well under SQLite's bound-parameter limit.
+fn move_events_to_bucket(
+    conn: &Connection,
+    bucket_bid: Option<i64>,
+    event_ids: &[i64],
+) -> rusqlite::Result<()> {
+    const BATCH: usize = 500;
+    for chunk in event_ids.chunks(BATCH) {
+        let placeholders = (0..chunk.len())
+            .map(|i| format!("?{}", i + 2))
+            .collect::<Vec<_>>()
+            .join(", ");
+        let sql = format!("UPDATE events SET bucketrow = ?1 WHERE id IN ({placeholders})");
+        let mut stmt = conn.prepare_cached(&sql)?;
+        let mut values: Vec<&dyn ToSql> = Vec::with_capacity(chunk.len() + 1);
+        values.push(&bucket_bid);
+        for id in chunk {
+            values.push(id);
+        }
+        stmt.execute(values.as_slice())?;
+    }
+    Ok(())
+}
+
 impl DatastoreInstance {
     pub fn new(
         conn: &Connection,
@@ -1143,22 +1268,21 @@ impl DatastoreInstance {
                 // Move only events that do not overlap the destination or another
                 // legacy event. A single overlapping cutover heartbeat must not strand
                 // years of disjoint history in the legacy bucket (ActivityWatch/aw-android#243).
-                conn.execute(
-                    "UPDATE events SET bucketrow = ?1
-                     WHERE id IN (
-                         SELECT old_event.id FROM events AS old_event
-                         WHERE old_event.bucketrow = ?2
-                           AND NOT EXISTS (
-                               SELECT 1 FROM events AS other_event
-                               WHERE other_event.id != old_event.id
-                                 AND other_event.bucketrow IN (?1, ?2)
-                                 AND old_event.starttime < other_event.endtime
-                                 AND other_event.starttime < old_event.endtime
-                           )
-                     )",
-                    [new_bucket.bid, old_bucket.bid],
-                )
-                .map_err(|err| {
+                //
+                // Overlaps are found with one sorted scan over both buckets instead of a
+                // correlated `NOT EXISTS` subquery per legacy event: that subquery had no
+                // lower bound on `starttime`, so it scanned every earlier event again for
+                // each row (O(n^2)). With a couple of years of Android history it kept the
+                // single datastore worker busy for hours, which blanked the web UI and
+                // produced ANRs in every main-thread datastore call (aw-android#261).
+                let movable = movable_legacy_event_ids(conn, old_bucket.bid, new_bucket.bid)
+                    .map_err(|err| {
+                        DatastoreError::InternalError(format!(
+                            "Failed to find mergeable events in '{}': {err}",
+                            old_id
+                        ))
+                    })?;
+                move_events_to_bucket(conn, new_bucket.bid, &movable).map_err(|err| {
                     DatastoreError::InternalError(format!(
                         "Failed to merge bucket '{}' into '{}': {err}",
                         old_id, new_id

aw-datastore/tests/datastore.rs

@@ -198,6 +198,119 @@ mod datastore_tests {
         assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 2);
     }
 
+    #[test]
+    fn test_migrate_test_bucket_names_zero_duration_event_at_shared_start_stays() {
+        // Strict overlap semantics: a zero-duration legacy event that starts exactly
+        // where a destination event starts is inside that event's open interval.
+        let ds = Datastore::new_in_memory(false);
+        let old_id = "aw-watcher-android-test_phone";
+        let new_id = "aw-watcher-android_phone";
+        create_named_test_bucket(&ds, old_id);
+        create_named_test_bucket(&ds, new_id);
+        let now = Utc::now();
+        ds.insert_events(
+            old_id,
+            &[test_event(now + Duration::minutes(1), Duration::zero())],
+        )
+        .unwrap();
+        ds.insert_events(new_id, &[test_event(now, Duration::minutes(30))])
+            .unwrap();
+
+        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
+        assert_eq!(ds.get_events(old_id, None, None, None).unwrap().len(), 1);
+        assert_eq!(ds.get_events(new_id, None, None, None).unwrap().len(), 1);
+    }
+
+    #[test]
+    fn test_migrate_test_bucket_names_merges_large_history_quickly() {
+        // Regression for ActivityWatch/aw-android#261: the merge used a correlated
+        // overlap subquery per legacy event (O(n^2)), which took hours on a phone
+        // with a couple of years of history and wedged the datastore worker. The
+        // sorted sweep must handle six-figure histories in seconds.
+        let ds = Datastore::new_in_memory(false);
+        let old_id = "aw-watcher-android-test_phone";
+        let new_id = "aw-watcher-android_phone";
+        create_named_test_bucket(&ds, old_id);
+        create_named_test_bucket(&ds, new_id);
+        let legacy_count = 100_000;
+        let destination_count = 10_000;
+        let start = Utc::now() - Duration::days(400);
+        let step = Duration::minutes(5);
+        let legacy: Vec<Event> = (0..legacy_count)
+            .map(|i| test_event(start + step * i, step - Duration::seconds(1)))
+            .collect();
+        for chunk in legacy.chunks(10_000) {
+            ds.insert_events(old_id, chunk).unwrap();
+        }
+        let cutover = start + step * legacy_count;
+        // One overlapping cutover event must stay behind while the rest merges.
+        ds.insert_events(
+            old_id,
+            &[test_event(
+                cutover - Duration::minutes(2),
+                Duration::minutes(4),
+            )],
+        )
+        .unwrap();
+        let destination: Vec<Event> = (0..destination_count)
+            .map(|i| test_event(cutover + step * i, step - Duration::seconds(1)))
+            .collect();
+        ds.insert_events(new_id, &destination).unwrap();
+
+        let started = std::time::Instant::now();
+        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
+        let elapsed = started.elapsed();
+        assert!(
+            elapsed < std::time::Duration::from_secs(30),
+            "merge took {elapsed:?}; the overlap scan regressed to quadratic time"
+        );
+        assert_eq!(
+            ds.get_event_count(old_id, None, None).unwrap(),
+            2,
+            "only the overlapping cutover pair stays in the legacy bucket"
+        );
+        // Everything except the legacy event that the cutover overlaps moved over.
+        assert_eq!(
+            ds.get_event_count(new_id, None, None).unwrap(),
+            (legacy_count - 1 + destination_count) as i64,
+        );
+    }
+
+    #[test]
+    fn test_migrate_test_bucket_names_heavy_overlap_is_linear() {
+        // Stacked histories must not reintroduce a quadratic open-set scan.
+        // Each event starts 1s later and outlives every later start, so the
+        // open heap grows to n; scanning it per event would be ~n².
+        let ds = Datastore::new_in_memory(false);
+        let old_id = "aw-watcher-android-test_phone";
+        let new_id = "aw-watcher-android_phone";
+        create_named_test_bucket(&ds, old_id);
+        create_named_test_bucket(&ds, new_id);
+        let n = 20_000;
+        let start = Utc::now();
+        let span = Duration::seconds(n as i64);
+        let legacy: Vec<Event> = (0..n)
+            .map(|i| test_event(start + Duration::seconds(i as i64), span))
+            .collect();
+        for chunk in legacy.chunks(5_000) {
+            ds.insert_events(old_id, chunk).unwrap();
+        }
+
+        let started = std::time::Instant::now();
+        assert_eq!(ds.migrate_test_bucket_names().unwrap(), 0);
+        let elapsed = started.elapsed();
+        assert!(
+            elapsed < std::time::Duration::from_secs(10),
+            "heavy-overlap merge took {elapsed:?}; open-set scan is quadratic"
+        );
+        assert_eq!(
+            ds.get_event_count(old_id, None, None).unwrap(),
+            n as i64,
+            "every stacked event overlaps another, so none should move"
+        );
+        assert_eq!(ds.get_event_count(new_id, None, None).unwrap(), 0);
+    }
+
     #[test]
     fn test_migrate_test_bucket_names_merges_interleaved_non_overlapping_events() {
         let ds = Datastore::new_in_memory(false);