laburnum 1.17.1

// Copyright Two Neutron Stars Incorporated and contributors
// SPDX-License-Identifier: BlueOak-1.0.0

use {
  super::storage::{
    TestPartition,
    TestPartitions,
    TestRecordData,
    Test1Partition,
  },
  crate::{
    Ident,
    database::{
      Database,
      chunk::RecordWriter,
      PartitionKey,
      query::SortKeyCondition,
    },
    record::LaburnumRecord,
  },
  std::{
    sync::{
      Arc,
      atomic::{AtomicBool, Ordering},
    },
    task::{Wake, Waker},
  },
};

struct TestWake(AtomicBool);

impl TestWake {
  fn new() -> Arc<Self> {
    Arc::new(Self(AtomicBool::new(false)))
  }

  fn was_woken(&self) -> bool {
    self.0.load(Ordering::SeqCst)
  }
}

impl Wake for TestWake {
  fn wake(self: Arc<Self>) {
    self.0.store(true, Ordering::SeqCst);
  }
}

fn make_test_record(value: &str) -> TestRecordData {
  TestRecordData::Laburnum(LaburnumRecord::WorkspaceConfig {
    value: value.to_string(),
  })
}

fn source_cache() -> crate::source::cache::reporter::SourceCacheReader {
  crate::source::cache::reporter::SourceCacheReader::new_empty_for_test()
}

#[test]
fn test_waker_exact_match_wakes() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  let tw = TestWake::new();
  let waker = Waker::from(tw.clone());

  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::Exact("sk1".to_string()),
    waker,
  );

  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("sk1", make_test_record("v1"));
  db.commit_chunk(writer.build(), &sc);

  assert!(tw.was_woken());
}

#[test]
fn test_waker_exact_mismatch_survives() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  let tw = TestWake::new();
  let waker = Waker::from(tw.clone());

  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::Exact("sk1".to_string()),
    waker,
  );

  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("sk2", make_test_record("v2"));
  db.commit_chunk(writer.build(), &sc);

  assert!(!tw.was_woken());

  // Waker should still be registered
  let wakers = db.partition_wakers.get(&TestPartition::KEY);
  assert!(wakers.is_some());
  assert_eq!(wakers.as_ref().map(|w| w.len()), Some(1));
}

#[test]
fn test_waker_begins_with_match() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  let tw = TestWake::new();
  let waker = Waker::from(tw.clone());

  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::BeginsWith("prefix".to_string()),
    waker,
  );

  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("prefix_foo", make_test_record("v1"));
  db.commit_chunk(writer.build(), &sc);

  assert!(tw.was_woken());
}

#[test]
fn test_waker_all_always_wakes() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  let tw = TestWake::new();
  let waker = Waker::from(tw.clone());

  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::All,
    waker,
  );

  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("anything", make_test_record("v1"));
  db.commit_chunk(writer.build(), &sc);

  assert!(tw.was_woken());
}

#[test]
fn test_waker_deletion_wakes() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  // First insert a record
  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("sk1", make_test_record("v1"));
  db.commit_chunk(writer.build(), &sc);

  // Register a waker for that sort key
  let tw = TestWake::new();
  let waker = Waker::from(tw.clone());

  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::Exact("sk1".to_string()),
    waker,
  );

  // Commit a chunk that clears the prefix containing "sk1"
  let mut writer = RecordWriter::new(Ident::new("task2"));
  writer.clear_prefix(TestPartition::KEY, "sk1");
  db.commit_chunk(writer.build(), &sc);

  assert!(tw.was_woken());
}

#[test]
fn test_waker_mixed_selective() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  let tw1 = TestWake::new();
  let tw2 = TestWake::new();

  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::Exact("a".to_string()),
    Waker::from(tw1.clone()),
  );
  db.register_partition_waker(
    TestPartition::KEY,
    SortKeyCondition::Exact("b".to_string()),
    Waker::from(tw2.clone()),
  );

  // Only commit "a"
  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("a", make_test_record("v1"));
  db.commit_chunk(writer.build(), &sc);

  assert!(tw1.was_woken());
  assert!(!tw2.was_woken());

  // Waker for "b" should still be registered
  let wakers = db.partition_wakers.get(&TestPartition::KEY);
  assert!(wakers.is_some());
  assert_eq!(wakers.as_ref().map(|w| w.len()), Some(1));
}

#[test]
fn test_waker_unaffected_partition_untouched() {
  let db = Database::<TestPartitions>::new();
  let sc = source_cache();

  let tw = TestWake::new();
  let waker = Waker::from(tw.clone());

  // Register waker on Test1Partition
  db.register_partition_waker(
    Test1Partition::KEY,
    SortKeyCondition::All,
    waker,
  );

  // Commit to TestPartition (different partition)
  let mut writer = RecordWriter::new(Ident::new("task1"));
  writer.insert::<TestPartition, _>("sk1", make_test_record("v1"));
  db.commit_chunk(writer.build(), &sc);

  assert!(!tw.was_woken());

  // Waker should still be registered on Test1Partition
  let wakers = db.partition_wakers.get(&Test1Partition::KEY);
  assert!(wakers.is_some());
  assert_eq!(wakers.as_ref().map(|w| w.len()), Some(1));
}