Phase 1: Windows log collection + full-text search
Extends the agent, ingest, storage, api, and web with Windows Event Log/ETW sourcing and Tantivy-backed free-text search, per the approved Phase 1 plan. - CLAUDE.md: materialized on disk (never existed as a file before) with a new Phase 1 "done looks like" section. - agent: Windows Event Log (EvtSubscribe) and ETW sources, Windows service wrapper (install/uninstall/run-service), both feature- and target_os-gated so Linux builds/tests/clippy stay unaffected. Also fixed two pre-existing Phase 0 clippy gaps (dead-code on default-features-only builds, a type-inference edge case) found while testing every feature combination properly for the first time. UNVERIFIED on real Windows -- no Windows toolchain existed anywhere in the build environment; flagged prominently in three places. - proto/ingest: new record_id field, assigned once server-side in ingest's gRPC front end so ClickHouse and Tantivy agree on the same ID for the same record. - storage: record_id column + bloom filter index, verified against a live ClickHouse. - search: new service, Tantivy index, rskafka consumer as an independent second consumer group on the same Redpanda topic ingest already reads. - api/web: new /search endpoint and page, sharing the query page's result-table shape and component. - hack/windows-fixture: sends realistic Windows-shaped data straight to ingest, so the pipeline's handling of it is verifiable without a Windows host. Verified end-to-end on the live docker-compose stack: the same record_id comes back from both /query and /search for the same log line, including for windows-fixture's synthetic Windows Event Log data. Real bugs found and fixed along the way: api/Dockerfile missing proto/ in its build context, search's logs being completely silent (RUST_LOG gap), and search/target/ missing from .gitignore/.dockerignore.
This commit is contained in:
@@ -0,0 +1,44 @@
|
||||
use anyhow::{Context, Result};
|
||||
use std::path::PathBuf;
|
||||
use std::time::Duration;
|
||||
|
||||
/// All via environment variables, same convention as /ingest and /api —
|
||||
/// no config file format for this service either.
|
||||
pub struct Config {
|
||||
pub grpc_listen_addr: String,
|
||||
pub redpanda_brokers: Vec<String>,
|
||||
pub redpanda_topic: String,
|
||||
pub index_path: PathBuf,
|
||||
pub offsets_path: PathBuf,
|
||||
pub commit_interval: Duration,
|
||||
}
|
||||
|
||||
impl Config {
|
||||
pub fn load() -> Result<Self> {
|
||||
let commit_interval_ms: u64 = getenv("COMMIT_INTERVAL_MS", "2000")
|
||||
.parse()
|
||||
.context("COMMIT_INTERVAL_MS must be a number")?;
|
||||
|
||||
Ok(Self {
|
||||
// Rust's SocketAddr parser needs a full address, unlike Go's
|
||||
// net package (ingest/api's ":PORT" convention won't parse
|
||||
// here).
|
||||
grpc_listen_addr: getenv("GRPC_LISTEN_ADDR", "0.0.0.0:50052"),
|
||||
redpanda_brokers: getenv("REDPANDA_BROKERS", "localhost:9092")
|
||||
.split(',')
|
||||
.map(str::to_string)
|
||||
.collect(),
|
||||
redpanda_topic: getenv("REDPANDA_TOPIC", "sentry.logs.raw"),
|
||||
index_path: PathBuf::from(getenv("INDEX_PATH", "/var/lib/sentry-search/index")),
|
||||
offsets_path: PathBuf::from(getenv(
|
||||
"OFFSETS_PATH",
|
||||
"/var/lib/sentry-search/offsets.json",
|
||||
)),
|
||||
commit_interval: Duration::from_millis(commit_interval_ms),
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
fn getenv(key: &str, fallback: &str) -> String {
|
||||
std::env::var(key).unwrap_or_else(|_| fallback.to_string())
|
||||
}
|
||||
@@ -0,0 +1,134 @@
|
||||
use anyhow::{Context, Result};
|
||||
use prost::Message;
|
||||
use rskafka::client::partition::UnknownTopicHandling;
|
||||
use rskafka::client::ClientBuilder;
|
||||
use std::sync::Arc;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
use crate::config::Config;
|
||||
use crate::index::SearchIndex;
|
||||
use crate::logsv1;
|
||||
use crate::offsets::OffsetStore;
|
||||
|
||||
/// Reads the same `sentry.logs.raw` topic ingest's ClickHouse-writer
|
||||
/// consumer reads, as an independent consumer group in spirit (its own
|
||||
/// offset tracking, own failure domain) even though rskafka doesn't speak
|
||||
/// Kafka's broker-side consumer-group protocol -- see offsets.rs. One
|
||||
/// task per partition; partition count comes from config rather than
|
||||
/// discovered dynamically, since it has to match what
|
||||
/// /transport/provision-topics.sh actually created anyway (documented
|
||||
/// cross-component contract, same as the topic name already is).
|
||||
pub async fn run(cfg: Arc<Config>, index: Arc<SearchIndex>, partition_count: i32) -> Result<()> {
|
||||
let client = ClientBuilder::new(cfg.redpanda_brokers.clone())
|
||||
.build()
|
||||
.await
|
||||
.context("building rskafka client")?;
|
||||
let client = Arc::new(client);
|
||||
|
||||
let offsets = OffsetStore::load(&cfg.offsets_path)
|
||||
.await
|
||||
.context("loading offset store")?;
|
||||
let offsets = Arc::new(Mutex::new(offsets));
|
||||
|
||||
// Periodic Tantivy commit, batched for throughput the same way
|
||||
// ingest's ClickHouse writer batches inserts rather than inserting
|
||||
// per-record.
|
||||
let commit_interval = cfg.commit_interval;
|
||||
let index_for_commit = Arc::clone(&index);
|
||||
tokio::spawn(async move {
|
||||
let mut ticker = tokio::time::interval(commit_interval);
|
||||
loop {
|
||||
ticker.tick().await;
|
||||
if let Err(e) = index_for_commit.commit().await {
|
||||
tracing::error!(error = %e, "periodic tantivy commit failed");
|
||||
}
|
||||
}
|
||||
});
|
||||
|
||||
let mut handles = Vec::with_capacity(partition_count as usize);
|
||||
for partition in 0..partition_count {
|
||||
let start_offset = offsets.lock().await.get(partition);
|
||||
let client = Arc::clone(&client);
|
||||
let index = Arc::clone(&index);
|
||||
let offsets = Arc::clone(&offsets);
|
||||
let topic = cfg.redpanda_topic.clone();
|
||||
handles.push(tokio::spawn(async move {
|
||||
consume_partition(client, topic, partition, start_offset, index, offsets).await
|
||||
}));
|
||||
}
|
||||
|
||||
for handle in handles {
|
||||
handle
|
||||
.await
|
||||
.context("partition consumer task panicked")??;
|
||||
}
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
async fn consume_partition(
|
||||
client: Arc<rskafka::client::Client>,
|
||||
topic: String,
|
||||
partition: i32,
|
||||
start_offset: i64,
|
||||
index: Arc<SearchIndex>,
|
||||
offsets: Arc<Mutex<OffsetStore>>,
|
||||
) -> Result<()> {
|
||||
let partition_client = client
|
||||
.partition_client(topic.clone(), partition, UnknownTopicHandling::Error)
|
||||
.await
|
||||
.with_context(|| format!("creating partition client for {topic}[{partition}]"))?;
|
||||
|
||||
let mut offset = start_offset;
|
||||
loop {
|
||||
let (records, _high_watermark) = partition_client
|
||||
.fetch_records(offset, 1..1_000_000, 5_000)
|
||||
.await
|
||||
.with_context(|| {
|
||||
format!("fetching records from {topic}[{partition}] at offset {offset}")
|
||||
})?;
|
||||
|
||||
if records.is_empty() {
|
||||
continue;
|
||||
}
|
||||
|
||||
for record_and_offset in &records {
|
||||
offset = record_and_offset.offset + 1;
|
||||
|
||||
let Some(value) = &record_and_offset.record.value else {
|
||||
continue;
|
||||
};
|
||||
let rec = match logsv1::LogRecord::decode(value.as_slice()) {
|
||||
Ok(rec) => rec,
|
||||
Err(e) => {
|
||||
tracing::warn!(error = %e, partition, "skipping unparseable message");
|
||||
continue;
|
||||
}
|
||||
};
|
||||
|
||||
if rec.record_id.is_empty() {
|
||||
// Shouldn't happen -- ingest's gRPC front end always
|
||||
// assigns this before producing -- but a message with no
|
||||
// ID can't be joined back to a ClickHouse row, so it's
|
||||
// useless to index.
|
||||
tracing::warn!(partition, "skipping record with empty record_id");
|
||||
continue;
|
||||
}
|
||||
|
||||
if let Err(e) = index.upsert(&rec.record_id, &rec.message).await {
|
||||
tracing::error!(error = %e, record_id = %rec.record_id, "failed to index record");
|
||||
}
|
||||
}
|
||||
|
||||
// Persisted after each fetched batch, not per-record: worst-case
|
||||
// reprocessing on an unclean restart is one batch, which
|
||||
// `SearchIndex::upsert`'s delete-then-add makes harmless anyway.
|
||||
{
|
||||
let mut offsets = offsets.lock().await;
|
||||
offsets.set(partition, offset);
|
||||
if let Err(e) = offsets.persist().await {
|
||||
tracing::error!(error = %e, partition, "failed to persist offset");
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,48 @@
|
||||
use std::sync::Arc;
|
||||
use tonic::{Request, Response, Status};
|
||||
|
||||
use crate::index::SearchIndex;
|
||||
use crate::searchv1;
|
||||
|
||||
const DEFAULT_LIMIT: usize = 100;
|
||||
|
||||
pub struct SearchServer {
|
||||
index: Arc<SearchIndex>,
|
||||
}
|
||||
|
||||
impl SearchServer {
|
||||
pub fn new(index: Arc<SearchIndex>) -> Self {
|
||||
Self { index }
|
||||
}
|
||||
}
|
||||
|
||||
#[tonic::async_trait]
|
||||
impl searchv1::search_service_server::SearchService for SearchServer {
|
||||
async fn search(
|
||||
&self,
|
||||
request: Request<searchv1::SearchRequest>,
|
||||
) -> Result<Response<searchv1::SearchResponse>, Status> {
|
||||
let req = request.into_inner();
|
||||
if req.query.trim().is_empty() {
|
||||
return Err(Status::invalid_argument("query must not be empty"));
|
||||
}
|
||||
|
||||
let limit = if req.limit == 0 {
|
||||
DEFAULT_LIMIT
|
||||
} else {
|
||||
req.limit as usize
|
||||
};
|
||||
|
||||
let index = Arc::clone(&self.index);
|
||||
let query = req.query.clone();
|
||||
// Tantivy's searcher is synchronous; run it on a blocking thread
|
||||
// so it doesn't stall the async runtime alongside the consumer
|
||||
// tasks.
|
||||
let record_ids = tokio::task::spawn_blocking(move || index.search(&query, limit))
|
||||
.await
|
||||
.map_err(|e| Status::internal(format!("search task panicked: {e}")))?
|
||||
.map_err(|e| Status::invalid_argument(format!("search failed: {e}")))?;
|
||||
|
||||
Ok(Response::new(searchv1::SearchResponse { record_ids }))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,192 @@
|
||||
use anyhow::{Context, Result};
|
||||
use std::path::Path;
|
||||
use tantivy::collector::TopDocs;
|
||||
use tantivy::query::QueryParser;
|
||||
use tantivy::schema::{Schema, Value, STORED, STRING, TEXT};
|
||||
use tantivy::{doc, Index, IndexReader, IndexWriter, ReloadPolicy, TantivyDocument, Term};
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
/// Minimal Tantivy index: a stable `record_id` (stored, exact-match) and
|
||||
/// tokenized `message` text. Everything else (timestamp, host, service,
|
||||
/// severity) is fetched by joining `record_id` back against ClickHouse in
|
||||
/// `/api`'s search handler, not duplicated in here — this stays a pure
|
||||
/// text index, not a second copy of the row.
|
||||
pub struct SearchIndex {
|
||||
index: Index,
|
||||
writer: Mutex<IndexWriter>,
|
||||
reader: IndexReader,
|
||||
record_id_field: tantivy::schema::Field,
|
||||
message_field: tantivy::schema::Field,
|
||||
}
|
||||
|
||||
/// 50MB is Tantivy's own suggested minimum writer heap budget; Phase 1
|
||||
/// has no real sizing data yet to tune this against.
|
||||
const WRITER_HEAP_BYTES: usize = 50_000_000;
|
||||
|
||||
impl SearchIndex {
|
||||
pub fn open_or_create(path: &Path) -> Result<Self> {
|
||||
std::fs::create_dir_all(path).context("creating tantivy index directory")?;
|
||||
|
||||
let mut schema_builder = Schema::builder();
|
||||
let record_id_field = schema_builder.add_text_field("record_id", STRING | STORED);
|
||||
let message_field = schema_builder.add_text_field("message", TEXT);
|
||||
let schema = schema_builder.build();
|
||||
|
||||
let dir = tantivy::directory::MmapDirectory::open(path)
|
||||
.context("opening tantivy mmap directory")?;
|
||||
let index =
|
||||
Index::open_or_create(dir, schema).context("opening/creating tantivy index")?;
|
||||
|
||||
let writer = index
|
||||
.writer(WRITER_HEAP_BYTES)
|
||||
.context("creating tantivy index writer")?;
|
||||
|
||||
let reader = index
|
||||
.reader_builder()
|
||||
.reload_policy(ReloadPolicy::OnCommitWithDelay)
|
||||
.try_into()
|
||||
.context("building tantivy index reader")?;
|
||||
|
||||
Ok(Self {
|
||||
index,
|
||||
writer: Mutex::new(writer),
|
||||
reader,
|
||||
record_id_field,
|
||||
message_field,
|
||||
})
|
||||
}
|
||||
|
||||
/// Upserts one record: delete-then-add on record_id. Tantivy segments
|
||||
/// are immutable, so this delete-then-add is the standard idiom for
|
||||
/// updates, not a workaround -- and it matters here specifically
|
||||
/// because /search's offset tracking is best-effort (see
|
||||
/// consumer.rs's OffsetStore), so the same record can genuinely be
|
||||
/// reprocessed after an unclean shutdown. Without this, that would
|
||||
/// silently duplicate documents instead of just re-writing the same
|
||||
/// one.
|
||||
pub async fn upsert(&self, record_id: &str, message: &str) -> Result<()> {
|
||||
let writer = self.writer.lock().await;
|
||||
let term = Term::from_field_text(self.record_id_field, record_id);
|
||||
writer.delete_term(term);
|
||||
writer
|
||||
.add_document(doc!(
|
||||
self.record_id_field => record_id,
|
||||
self.message_field => message,
|
||||
))
|
||||
.context("adding document to tantivy index")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn commit(&self) -> Result<()> {
|
||||
let mut writer = self.writer.lock().await;
|
||||
writer.commit().context("committing tantivy index")?;
|
||||
// Explicit reload rather than relying solely on ReloadPolicy::
|
||||
// OnCommitWithDelay's background timing: callers of `commit()`
|
||||
// (the periodic ticker in consumer.rs, and tests) expect a
|
||||
// committed document to be immediately searchable, not visible
|
||||
// after some undocumented delay.
|
||||
self.reader.reload().context("reloading tantivy reader after commit")?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
/// Runs a Tantivy query-parser query against the `message` field,
|
||||
/// returning matching record_ids, most-relevant first. Phase 1: no
|
||||
/// pagination, no score exposed to the caller — just IDs for `/api`
|
||||
/// to join against ClickHouse.
|
||||
pub fn search(&self, query: &str, limit: usize) -> Result<Vec<String>> {
|
||||
let searcher = self.reader.searcher();
|
||||
let query_parser = QueryParser::for_index(&self.index, vec![self.message_field]);
|
||||
let parsed_query = query_parser
|
||||
.parse_query(query)
|
||||
.context("parsing search query")?;
|
||||
let top_docs = searcher
|
||||
.search(&parsed_query, &TopDocs::with_limit(limit))
|
||||
.context("executing search")?;
|
||||
|
||||
let mut ids = Vec::with_capacity(top_docs.len());
|
||||
for (_score, doc_address) in top_docs {
|
||||
let retrieved: TantivyDocument = searcher
|
||||
.doc(doc_address)
|
||||
.context("retrieving matched document")?;
|
||||
if let Some(value) = retrieved.get_first(self.record_id_field) {
|
||||
if let Some(s) = value.as_str() {
|
||||
ids.push(s.to_string());
|
||||
}
|
||||
}
|
||||
}
|
||||
Ok(ids)
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
fn new_test_index() -> (SearchIndex, tempfile::TempDir) {
|
||||
let dir = tempfile::tempdir().expect("creating temp dir");
|
||||
let index = SearchIndex::open_or_create(dir.path()).expect("opening tantivy index");
|
||||
(index, dir)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_and_search_finds_matching_message() {
|
||||
let (index, _dir) = new_test_index();
|
||||
index.upsert("id-1", "hello world").await.unwrap();
|
||||
index.upsert("id-2", "goodbye moon").await.unwrap();
|
||||
index.commit().await.unwrap();
|
||||
|
||||
let results = index.search("hello", 10).unwrap();
|
||||
assert_eq!(results, vec!["id-1".to_string()]);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_before_commit_finds_nothing() {
|
||||
let (index, _dir) = new_test_index();
|
||||
index.upsert("id-1", "hello world").await.unwrap();
|
||||
// no commit yet
|
||||
let results = index.search("hello", 10).unwrap();
|
||||
assert!(results.is_empty(), "expected no results before commit, got {results:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn upsert_same_id_twice_does_not_duplicate() {
|
||||
let (index, _dir) = new_test_index();
|
||||
index.upsert("id-1", "hello world").await.unwrap();
|
||||
index.commit().await.unwrap();
|
||||
index.upsert("id-1", "hello world again").await.unwrap();
|
||||
index.commit().await.unwrap();
|
||||
|
||||
let results = index.search("hello", 10).unwrap();
|
||||
assert_eq!(
|
||||
results.len(),
|
||||
1,
|
||||
"expected exactly one result after re-upserting the same record_id, got {results:?}"
|
||||
);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_respects_limit() {
|
||||
let (index, _dir) = new_test_index();
|
||||
for i in 0..5 {
|
||||
index
|
||||
.upsert(&format!("id-{i}"), "shared term")
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
index.commit().await.unwrap();
|
||||
|
||||
let results = index.search("shared", 2).unwrap();
|
||||
assert_eq!(results.len(), 2);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn search_supports_phrase_queries() {
|
||||
let (index, _dir) = new_test_index();
|
||||
index.upsert("id-1", "the quick brown fox").await.unwrap();
|
||||
index.upsert("id-2", "quick and brown but not adjacent fox").await.unwrap();
|
||||
index.commit().await.unwrap();
|
||||
|
||||
let results = index.search("\"quick brown\"", 10).unwrap();
|
||||
assert_eq!(results, vec!["id-1".to_string()]);
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,67 @@
|
||||
mod config;
|
||||
mod consumer;
|
||||
mod grpc;
|
||||
mod index;
|
||||
mod offsets;
|
||||
|
||||
pub mod logsv1 {
|
||||
tonic::include_proto!("sentry.logs.v1");
|
||||
}
|
||||
pub mod searchv1 {
|
||||
tonic::include_proto!("sentry.search.v1");
|
||||
}
|
||||
|
||||
use anyhow::{Context, Result};
|
||||
use config::Config;
|
||||
use index::SearchIndex;
|
||||
use std::sync::Arc;
|
||||
use tonic::transport::Server;
|
||||
|
||||
/// Matches /transport/provision-topics.sh's default
|
||||
/// REDPANDA_TOPIC_PARTITIONS -- documented cross-component contract, not
|
||||
/// discovered dynamically. See consumer.rs.
|
||||
const DEFAULT_PARTITION_COUNT: i32 = 6;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<()> {
|
||||
tracing_subscriber::fmt()
|
||||
.with_env_filter(tracing_subscriber::EnvFilter::from_default_env())
|
||||
.init();
|
||||
|
||||
let cfg = Arc::new(Config::load().context("loading config")?);
|
||||
|
||||
let index = Arc::new(
|
||||
SearchIndex::open_or_create(&cfg.index_path).context("opening tantivy index")?,
|
||||
);
|
||||
|
||||
let partition_count: i32 = std::env::var("REDPANDA_TOPIC_PARTITIONS")
|
||||
.ok()
|
||||
.and_then(|s| s.parse().ok())
|
||||
.unwrap_or(DEFAULT_PARTITION_COUNT);
|
||||
|
||||
let consumer_cfg = Arc::clone(&cfg);
|
||||
let consumer_index = Arc::clone(&index);
|
||||
let consumer_handle = tokio::spawn(async move {
|
||||
if let Err(e) = consumer::run(consumer_cfg, consumer_index, partition_count).await {
|
||||
tracing::error!(error = %e, "redpanda consumer exited with error");
|
||||
}
|
||||
});
|
||||
|
||||
let addr = cfg
|
||||
.grpc_listen_addr
|
||||
.parse()
|
||||
.context("parsing GRPC_LISTEN_ADDR")?;
|
||||
tracing::info!(addr = %cfg.grpc_listen_addr, "search gRPC server listening");
|
||||
|
||||
let search_server = grpc::SearchServer::new(Arc::clone(&index));
|
||||
Server::builder()
|
||||
.add_service(searchv1::search_service_server::SearchServiceServer::new(
|
||||
search_server,
|
||||
))
|
||||
.serve(addr)
|
||||
.await
|
||||
.context("gRPC server failed")?;
|
||||
|
||||
consumer_handle.abort();
|
||||
Ok(())
|
||||
}
|
||||
@@ -0,0 +1,93 @@
|
||||
use anyhow::{Context, Result};
|
||||
use std::collections::HashMap;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
/// Tracks per-partition offsets in a plain JSON file next to the Tantivy
|
||||
/// index, since rskafka is a low-level client with no built-in consumer-
|
||||
/// group coordination/offset-commit protocol (unlike kafka-go on the
|
||||
/// ingest side) -- there's no broker-side group to commit to here, so
|
||||
/// this service owns its own offset bookkeeping.
|
||||
///
|
||||
/// Best-effort, not exactly-once: if the process dies between processing
|
||||
/// a record and persisting its offset, that record gets reprocessed on
|
||||
/// restart. This is fine because `SearchIndex::upsert` is delete-then-add
|
||||
/// on `record_id` -- reprocessing the same record overwrites the same
|
||||
/// document rather than duplicating it.
|
||||
#[derive(Debug)]
|
||||
pub struct OffsetStore {
|
||||
path: PathBuf,
|
||||
offsets: HashMap<i32, i64>,
|
||||
}
|
||||
|
||||
impl OffsetStore {
|
||||
pub async fn load(path: &Path) -> Result<Self> {
|
||||
let offsets = match tokio::fs::read(path).await {
|
||||
Ok(bytes) => serde_json::from_slice(&bytes)
|
||||
.with_context(|| format!("parsing offsets file {}", path.display()))?,
|
||||
Err(e) if e.kind() == std::io::ErrorKind::NotFound => HashMap::new(),
|
||||
Err(e) => return Err(e).with_context(|| format!("reading offsets file {}", path.display())),
|
||||
};
|
||||
Ok(Self {
|
||||
path: path.to_path_buf(),
|
||||
offsets,
|
||||
})
|
||||
}
|
||||
|
||||
/// Next offset to fetch for a partition -- 0 (earliest) if never
|
||||
/// recorded before.
|
||||
pub fn get(&self, partition: i32) -> i64 {
|
||||
self.offsets.get(&partition).copied().unwrap_or(0)
|
||||
}
|
||||
|
||||
pub fn set(&mut self, partition: i32, offset: i64) {
|
||||
self.offsets.insert(partition, offset);
|
||||
}
|
||||
|
||||
pub async fn persist(&self) -> Result<()> {
|
||||
if let Some(parent) = self.path.parent() {
|
||||
tokio::fs::create_dir_all(parent)
|
||||
.await
|
||||
.with_context(|| format!("creating offsets directory {}", parent.display()))?;
|
||||
}
|
||||
let bytes = serde_json::to_vec_pretty(&self.offsets).context("serializing offsets")?;
|
||||
tokio::fs::write(&self.path, bytes)
|
||||
.await
|
||||
.with_context(|| format!("writing offsets file {}", self.path.display()))
|
||||
}
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
|
||||
#[tokio::test]
|
||||
async fn get_defaults_to_zero_for_unknown_partition() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let store = OffsetStore::load(&dir.path().join("offsets.json")).await.unwrap();
|
||||
assert_eq!(store.get(0), 0);
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn missing_file_loads_as_empty_not_an_error() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("does-not-exist.json");
|
||||
let store = OffsetStore::load(&path).await;
|
||||
assert!(store.is_ok(), "expected a missing offsets file to load as empty, got {store:?}");
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn persist_and_reload_round_trips() {
|
||||
let dir = tempfile::tempdir().unwrap();
|
||||
let path = dir.path().join("offsets.json");
|
||||
|
||||
let mut store = OffsetStore::load(&path).await.unwrap();
|
||||
store.set(0, 42);
|
||||
store.set(1, 7);
|
||||
store.persist().await.unwrap();
|
||||
|
||||
let reloaded = OffsetStore::load(&path).await.unwrap();
|
||||
assert_eq!(reloaded.get(0), 42);
|
||||
assert_eq!(reloaded.get(1), 7);
|
||||
assert_eq!(reloaded.get(2), 0, "unset partition should still default to 0");
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user