use std::{
path::{Path, PathBuf},
time::{Duration, Instant},
};
use serde_json::Value;
use crate::{
index::DocIndexEntry,
schema::SchemaRegistry,
txlog::{JournalEntry, TransactionLog},
Error, IndexRegistry, Layout, Result,
};
#[derive(Debug)]
enum WriteOp {
Put { doc_id: String, doc: Value },
Delete { doc_id: String },
}
#[derive(Debug, Clone)]
pub struct WriteConfig {
pub max_pending_ops: usize,
pub max_pending_bytes: usize,
pub flush_interval_ms: u64,
}
impl Default for WriteConfig {
fn default() -> Self {
WriteConfig {
max_pending_ops: 1000,
max_pending_bytes: 10 * 1024 * 1024, flush_interval_ms: 1000,
}
}
}
pub struct BufferedWriter {
root: PathBuf,
collection: String,
config: WriteConfig,
pending_ops: Vec<WriteOp>,
pending_bytes: usize,
last_flush: Instant,
journal: TransactionLog,
index: IndexRegistry,
schema_registry: SchemaRegistry,
}
impl BufferedWriter {
pub fn new(
root: impl AsRef<Path>,
collection: impl AsRef<str>,
config: WriteConfig,
) -> Result<Self> {
let root = root.as_ref().to_path_buf();
let collection = collection.as_ref().to_string();
Layout::init_collection(&root, &collection)?;
let index = IndexRegistry::load(&root, &collection).unwrap_or_default();
let schema_registry =
SchemaRegistry::load(&root, &collection).unwrap_or_else(|_| SchemaRegistry::new(false));
let journal = TransactionLog::open(&root, &collection)?;
Ok(BufferedWriter {
root,
collection,
config,
pending_ops: Vec::new(),
pending_bytes: 0,
last_flush: Instant::now(),
journal,
index,
schema_registry,
})
}
pub fn put(&mut self, doc_id: impl Into<String>, doc: Value) -> Result<()> {
let doc_id = doc_id.into();
Layout::validate_doc_id(&doc_id)?;
let doc_size = serde_json::to_string(&doc)?.len();
self.pending_bytes += doc_size;
self.pending_ops.push(WriteOp::Put { doc_id, doc });
self.maybe_flush()?;
Ok(())
}
pub fn delete(&mut self, doc_id: impl Into<String>) -> Result<()> {
let doc_id = doc_id.into();
self.pending_ops.push(WriteOp::Delete { doc_id });
self.maybe_flush()?;
Ok(())
}
fn maybe_flush(&mut self) -> Result<()> {
let should_flush = self.pending_ops.len() >= self.config.max_pending_ops
|| self.pending_bytes >= self.config.max_pending_bytes
|| self.last_flush.elapsed() >= Duration::from_millis(self.config.flush_interval_ms);
if should_flush {
self.flush()?;
}
Ok(())
}
pub fn flush(&mut self) -> Result<()> {
if self.pending_ops.is_empty() {
return Ok(());
}
let ops = std::mem::take(&mut self.pending_ops);
self.pending_bytes = 0;
self.last_flush = Instant::now();
for op in ops {
match op {
WriteOp::Put { doc_id, doc } => {
self.write_doc(&doc_id, &doc)?;
}
WriteOp::Delete { doc_id } => {
self.delete_doc(&doc_id)?;
}
}
}
self.journal.commit()?;
self.index.save(&self.root, &self.collection)?;
self.schema_registry.save(&self.root, &self.collection)?;
Ok(())
}
fn write_doc(&mut self, doc_id: &str, doc: &Value) -> Result<()> {
let schema_id = self.schema_registry.register(doc)?;
let docs_dir = Layout::docs_dir(&self.root, &self.collection);
let final_path = Layout::doc_file(&self.root, &self.collection, doc_id);
let tmp_path = docs_dir.join(format!(".{}.tmp", doc_id));
std::fs::create_dir_all(&docs_dir)?;
let content = serde_json::to_string_pretty(doc)?;
let size = content.len() as u64;
std::fs::write(&tmp_path, &content)?;
self.journal
.append(&JournalEntry::put(doc_id, &schema_id, size))?;
std::fs::rename(&tmp_path, &final_path)?;
let mtime = std::fs::metadata(&final_path)
.ok()
.and_then(|m| m.modified().ok())
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map(|d| d.as_secs())
.unwrap_or(0);
self.index.put(DocIndexEntry {
doc_id: doc_id.to_string(),
schema_id,
size,
mtime,
});
Ok(())
}
fn delete_doc(&mut self, doc_id: &str) -> Result<()> {
let path = Layout::doc_file(&self.root, &self.collection, doc_id);
if !path.exists() {
return Err(Error::DocumentNotFound(doc_id.to_string()));
}
if let Some(entry) = self.index.get(doc_id) {
self.schema_registry.unregister(&entry.schema_id);
}
self.journal.append(&JournalEntry::delete(doc_id))?;
std::fs::remove_file(&path)?;
self.index.remove(doc_id);
Ok(())
}
pub fn len(&self) -> usize {
self.index.len()
}
pub fn is_empty(&self) -> bool {
self.index.is_empty()
}
}
impl Drop for BufferedWriter {
fn drop(&mut self) {
let _ = self.flush();
}
}
pub struct SyncWriter {
root: PathBuf,
collection: String,
journal: TransactionLog,
index: IndexRegistry,
schema_registry: SchemaRegistry,
}
impl SyncWriter {
pub fn new(root: impl AsRef<Path>, collection: impl AsRef<str>) -> Result<Self> {
let root = root.as_ref().to_path_buf();
let collection = collection.as_ref().to_string();
Layout::init_collection(&root, &collection)?;
let index = IndexRegistry::load(&root, &collection).unwrap_or_default();
let schema_registry =
SchemaRegistry::load(&root, &collection).unwrap_or_else(|_| SchemaRegistry::new(false));
let journal = TransactionLog::open(&root, &collection)?;
Ok(SyncWriter {
root,
collection,
journal,
index,
schema_registry,
})
}
pub fn put(&mut self, doc_id: &str, doc: &Value) -> Result<()> {
Layout::validate_doc_id(doc_id)?;
let schema_id = self.schema_registry.register(doc)?;
let docs_dir = Layout::docs_dir(&self.root, &self.collection);
let final_path = Layout::doc_file(&self.root, &self.collection, doc_id);
let tmp_path = docs_dir.join(format!(".{}.tmp", doc_id));
std::fs::create_dir_all(&docs_dir)?;
let content = serde_json::to_string_pretty(doc)?;
let size = content.len() as u64;
std::fs::write(&tmp_path, &content)?;
self.journal
.append(&JournalEntry::put(doc_id, &schema_id, size))?;
std::fs::rename(&tmp_path, &final_path)?;
let mtime = std::fs::metadata(&final_path)
.ok()
.and_then(|m| m.modified().ok())
.and_then(|t| t.duration_since(std::time::UNIX_EPOCH).ok())
.map(|d| d.as_secs())
.unwrap_or(0);
self.index.put(DocIndexEntry {
doc_id: doc_id.to_string(),
schema_id,
size,
mtime,
});
self.journal.commit()?;
self.index.save(&self.root, &self.collection)?;
self.schema_registry.save(&self.root, &self.collection)?;
Ok(())
}
pub fn delete(&mut self, doc_id: &str) -> Result<()> {
let path = Layout::doc_file(&self.root, &self.collection, doc_id);
if !path.exists() {
return Err(Error::DocumentNotFound(doc_id.to_string()));
}
if let Some(entry) = self.index.get(doc_id) {
self.schema_registry.unregister(&entry.schema_id);
}
self.journal.append(&JournalEntry::delete(doc_id))?;
std::fs::remove_file(&path)?;
self.index.remove(doc_id);
self.journal.commit()?;
self.index.save(&self.root, &self.collection)?;
self.schema_registry.save(&self.root, &self.collection)?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use serde_json::json;
use tempfile::TempDir;
use super::*;
#[test]
fn test_buffered_writer() {
let tmp = TempDir::new().unwrap();
let root = tmp.path();
Layout::init_root(root).unwrap();
let config = WriteConfig {
max_pending_ops: 2,
..Default::default()
};
let mut writer = BufferedWriter::new(root, "test", config).unwrap();
writer.put("doc1", json!({"name": "alice"})).unwrap();
writer.put("doc2", json!({"name": "bob"})).unwrap();
assert!(Layout::doc_file(root, "test", "doc1").exists());
assert!(Layout::doc_file(root, "test", "doc2").exists());
}
#[test]
fn test_sync_writer() {
let tmp = TempDir::new().unwrap();
let root = tmp.path();
Layout::init_root(root).unwrap();
let mut writer = SyncWriter::new(root, "test").unwrap();
writer.put("doc1", &json!({"name": "alice"})).unwrap();
assert!(Layout::doc_file(root, "test", "doc1").exists());
writer.delete("doc1").unwrap();
assert!(!Layout::doc_file(root, "test", "doc1").exists());
}
}