use super::super::*;
struct DrainerThreadExit {
pipe: Option<std::os::unix::net::UnixStream>,
finished: std::sync::mpsc::Sender<std::thread::ThreadId>,
dispatch: Option<tracing::dispatcher::DefaultGuard>,
}
impl Drop for DrainerThreadExit {
fn drop(&mut self) {
drop(self.pipe.take());
if let Err(error) = self.finished.send(std::thread::current().id()) {
eprintln!("drainer exit watch failed: {error}");
}
drop(self.dispatch.take());
}
}
thread_local! {
static DRAINER_THREAD_EXIT: std::cell::RefCell<Option<DrainerThreadExit>> =
const { std::cell::RefCell::new(None) };
}
struct DrainerProbe {
pipe: std::os::unix::net::UnixStream,
finished: std::sync::mpsc::Sender<std::thread::ThreadId>,
dispatch: tracing::Dispatch,
pause: Arc<AppendPause>,
panics: bool,
}
type RoutedBatchTrace = Arc<Mutex<Vec<(String, Vec<Vec<u8>>)>>>;
struct DrainerProbeStore {
inner: Arc<dyn liminal::durability::DurableStore>,
probe: Arc<Mutex<Option<DrainerProbe>>>,
routed_batches: RoutedBatchTrace,
}
impl std::fmt::Debug for DrainerProbeStore {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
formatter
.debug_struct("DrainerProbeStore")
.finish_non_exhaustive()
}
}
impl DrainerProbeStore {
fn boundary(&self, key: &str) -> Result<(), liminal::durability::DurabilityError> {
if !matches!(key, "liminal.events:0" | "liminal.cost:0") {
return Ok(());
}
let probe = self.probe.lock().map_err(append_pause_error)?.take();
if let Some(probe) = probe {
let dispatch = tracing::dispatcher::set_default(&probe.dispatch);
DRAINER_THREAD_EXIT.with(|slot| {
*slot.borrow_mut() = Some(DrainerThreadExit {
pipe: Some(probe.pipe),
finished: probe.finished,
dispatch: Some(dispatch),
});
});
probe.pause.wait_for_release()?;
if probe.panics {
std::panic::resume_unwind(Box::new("seeded record drainer panic".to_owned()));
}
}
Ok(())
}
}
#[async_trait::async_trait]
impl liminal::durability::DurableStore for DrainerProbeStore {
async fn apply_routed_batch(
&self,
route_key: &str,
items: Vec<liminal::durability::RoutedBatchItem>,
) -> Result<(), liminal::durability::DurabilityError> {
let keys = items.iter().map(|item| item.key.clone()).collect();
self.inner.apply_routed_batch(route_key, items).await?;
self.routed_batches
.lock()
.map_err(append_pause_error)?
.push((route_key.to_owned(), keys));
Ok(())
}
async fn get_routed(
&self,
route_key: &str,
key: &[u8],
) -> Result<Option<Vec<u8>>, liminal::durability::DurabilityError> {
self.inner.get_routed(route_key, key).await
}
async fn append(
&self,
key: &str,
payload: Vec<u8>,
expected: u64,
) -> Result<u64, liminal::durability::DurabilityError> {
self.boundary(key)?;
self.inner.append(key, payload, expected).await
}
async fn append_all(
&self,
key: &str,
payloads: Vec<Vec<u8>>,
expected: u64,
) -> Result<u64, liminal::durability::DurabilityError> {
self.boundary(key)?;
self.inner.append_all(key, payloads, expected).await
}
async fn read_from(
&self,
key: &str,
offset: u64,
limit: usize,
) -> Result<Vec<liminal::durability::StoredEntry>, liminal::durability::DurabilityError> {
self.inner.read_from(key, offset, limit).await
}
async fn read_at(
&self,
key: &str,
sequence: u64,
) -> Result<Option<liminal::durability::StoredEntry>, liminal::durability::DurabilityError>
{
self.inner.read_at(key, sequence).await
}
async fn cas(
&self,
key: &str,
old: u64,
new: u64,
) -> Result<(), liminal::durability::DurabilityError> {
self.inner.cas(key, old, new).await
}
async fn read_stream_next_seq(
&self,
stream_key: &str,
) -> Result<Option<u64>, liminal::durability::DurabilityError> {
self.inner.read_stream_next_seq(stream_key).await
}
async fn read_value(
&self,
key: &str,
) -> Result<Option<u64>, liminal::durability::DurabilityError> {
self.inner.read_value(key).await
}
async fn scan_page(
&self,
prefix: &str,
cursor: Option<liminal::durability::ScanCursor>,
limit: usize,
) -> Result<liminal::durability::ScanPage, liminal::durability::DurabilityError> {
self.inner.scan_page(prefix, cursor, limit).await
}
async fn flush(&self) -> Result<(), liminal::durability::DurabilityError> {
self.inner.flush().await
}
}
struct DrainerWatch {
pipe: std::os::unix::net::UnixStream,
finished: std::sync::mpsc::Receiver<std::thread::ThreadId>,
log: HeadLog,
}
impl DrainerWatch {
fn finish(mut self) -> Result<String, Box<dyn Error>> {
let thread = self.finished.recv()?;
assert_ne!(thread, std::thread::current().id());
self.pipe.set_nonblocking(true)?;
let mut byte = [0_u8];
let eof = std::io::Read::read(&mut self.pipe, &mut byte);
assert!(
matches!(eof, Ok(0)),
"drainer teardown left the owned exit watch open: {eof:?}"
);
drop(self.pipe);
Ok(String::from_utf8(
self.log.0.lock().map_err(test_sink_error)?.clone(),
)?)
}
}
pub(super) struct DrainerCase {
pub(super) fixture: Option<WireFixture>,
pub(super) queue: Arc<crate::server::self_streams::queue::RecordQueue>,
pub(super) pause: Arc<AppendPause>,
watch: Option<DrainerWatch>,
pub(super) store: Arc<dyn liminal::durability::DurableStore>,
pub(super) counts: liminal::durability::StoreCountsHandle,
pub(super) root: Option<tempfile::TempDir>,
routed_batches: RoutedBatchTrace,
}
impl DrainerCase {
pub(super) fn new(panics: bool) -> Result<Self, Box<dyn Error>> {
let (pipe, watched) = std::os::unix::net::UnixStream::pair()?;
let (finished, received) = std::sync::mpsc::channel();
let pause = Arc::new(AppendPause::default());
let log = HeadLog(Arc::new(Mutex::new(Vec::new())));
let writer = log.clone();
let subscriber = tracing_subscriber::fmt()
.without_time()
.with_ansi(false)
.with_max_level(tracing::Level::DEBUG)
.with_writer(move || writer.clone())
.finish();
let probe = Arc::new(Mutex::new(None));
let store_probe = Arc::clone(&probe);
let routed_batches = Arc::new(Mutex::new(Vec::new()));
let store_batches = Arc::clone(&routed_batches);
let fixture = WireFixture::with_store(move |inner| {
Arc::new(liminal::durability::CountingStore::new(DrainerProbeStore {
inner,
probe: store_probe,
routed_batches: store_batches,
}))
})?;
let queue = fixture.sink.queue_for_test();
let store = fixture.services.durable_store();
let counts = fixture.services.store_counts();
*probe.lock().map_err(test_sink_error)? = Some(DrainerProbe {
pipe,
finished,
dispatch: tracing::Dispatch::new(subscriber),
pause: Arc::clone(&pause),
panics,
});
Ok(Self {
fixture: Some(fixture),
queue,
pause,
watch: Some(DrainerWatch {
pipe: watched,
finished: received,
log,
}),
store,
counts,
root: None,
routed_batches,
})
}
pub(super) fn fixture(&self) -> Result<&WireFixture, Box<dyn Error>> {
self.fixture
.as_ref()
.ok_or_else(|| "drainer fixture already released".into())
}
pub(super) fn start_event(&self) -> Result<u64, Box<dyn Error>> {
let sequence = self
.fixture()?
.sink
.enqueue_event_for_test(ServerEvent::ConnectionOpened { connection: 1 })?;
self.pause.entered()?;
Ok(sequence)
}
pub(super) fn release_services(&mut self) -> Result<(), Box<dyn Error>> {
let WireFixture {
services,
supervisor,
sink,
wake,
root,
..
} = self
.fixture
.take()
.ok_or("drainer fixture already released")?;
supervisor.shutdown();
let batches = self.routed_batches.lock().map_err(test_sink_error)?;
assert_eq!(
batches.as_slice(),
&[(
"liminal:participant-conversation-registry".to_owned(),
vec![b"registry/capacity/v1".to_vec()],
)]
);
eprintln!("DRAINER_CLEAN_STOP_ROUTED_BATCHES {batches:?}");
drop(batches);
drop(supervisor);
drop(sink);
drop(services);
drop(wake);
self.root = Some(root);
Ok(())
}
pub(super) fn finish(&mut self) -> Result<String, Box<dyn Error>> {
self.watch
.take()
.ok_or("drainer exit watch already consumed")?
.finish()
}
}
impl Drop for DrainerCase {
fn drop(&mut self) {
if let Err(error) = self.pause.release() {
eprintln!("drainer probe release failed: {error}");
}
if let Some(fixture) = self.fixture.take() {
fixture.supervisor.shutdown();
drop(fixture);
}
}
}