insightful/

lib.rs

use parking_lot::RwLock;
use serde::Deserialize;
use std::{
    collections::BTreeMap,
    error::Error,
    fmt::{Debug, Display},
    fs::File,
    io::{BufRead, BufReader, Write},
    path::PathBuf,
    sync::Arc,
    time::Duration,
};
use tantivy::{
    Document, Index, IndexReader, IndexWriter, TantivyDocument,
    directory::MmapDirectory,
    query::QueryParser,
    schema::{DateOptions, INDEXED, STORED, STRING, Schema, TEXT},
};
use tokio::task::JoinHandle;
use tracing::instrument;

#[derive(Deserialize)]
struct LogSpan {
    // name: String,

    // env: Option<String>,
    domain: Option<String>,
    // id: Option<String>,

    // i: Option<u8>,
    // nm: Option<String>,

    // trigger: Option<String>,

    // flavor: Option<String>,
}

#[derive(Deserialize)]
struct LogFields {
    // message: Option<String>,

    // ip: Option<String>,
    // version: Option<String>,
    // method: Option<String>,
    // uri: Option<String>,
    // headers: Option<String>,
    // status: Option<u16>,

    // token: Option<String>,

    // ts: Option<String>,
    // path: Option<String>,
    // query: Option<String>,
    // correlation: Option<String>,
    // flags: Option<String>,
}

#[derive(Deserialize)]
struct LogLine {
    // timestamp: String,
    // level: String,
    // fields: LogFields,
    spans: Option<Vec<LogSpan>>,
}

#[derive(Deserialize, Clone)]
#[serde(rename_all = "camelCase")]
pub enum QueryFormat {
    All,
    Count,
    Top,
}

impl QueryFormat {
    fn format(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            &Self::All => write!(f, "all"),
            &Self::Count => write!(f, "count"),
            &Self::Top => write!(f, "top"),
        }
    }
}

impl Display for QueryFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.format(f)
    }
}

impl Debug for QueryFormat {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.format(f)
    }
}

pub struct Manager {
    dir: PathBuf,

    // todo: needs more than just static interval,
    // todo: but this gets *something* for now.
    index_interval_s: u64,
    max_ttl_days: u16,
    max_lines: u32,
    cleanup_interval_s: u64,

    domain_indexes: Arc<RwLock<BTreeMap<String, (Schema, Index, IndexReader, QueryParser)>>>,
}

impl Manager {
    pub fn new(
        dir: PathBuf,

        index_interval_s: u64,
        max_ttl_days: u16,
        max_lines: u32,
        cleanup_interval_s: u64,
    ) -> Self {
        Manager {
            dir,

            index_interval_s,
            max_ttl_days,
            max_lines,
            cleanup_interval_s,

            domain_indexes: Arc::new(RwLock::new(BTreeMap::new())),
        }
    }

    #[instrument(skip(self), err)]
    pub fn add_domain(&self, domain: String) -> Result<(), Box<dyn Error>> {
        let mut schema_builder = Schema::builder();

        let ts_opts = DateOptions::from(INDEXED)
            .set_stored()
            .set_fast()
            .set_precision(tantivy::schema::DateTimePrecision::Microseconds);

        // todo: don't store the values, just index them and store
        // todo: logs in separate file, compressed.
        //
        // ?? store file path and log line in the index
        let timestamp = schema_builder.add_date_field("timestamp", ts_opts);
        let level = schema_builder.add_text_field("level", STRING | STORED);
        let fields = schema_builder.add_json_field("fields", TEXT | STORED);
        let spans = schema_builder.add_json_field("spans", TEXT | STORED);

        let schema = schema_builder.build();

        let index_dir = self.dir.join("search").join(&domain);
        std::fs::create_dir_all(&index_dir)?;

        let dir = MmapDirectory::open(index_dir)?;
        let index = Index::open_or_create(dir, schema.clone())?;

        let reader = index
            .reader_builder()
            .reload_policy(tantivy::ReloadPolicy::OnCommitWithDelay);

        let query_parser = QueryParser::for_index(&index, vec![timestamp, level, fields, spans]);

        let mut di = self.domain_indexes.write();
        di.insert(domain, (schema, index, reader.try_into()?, query_parser));

        Ok(())
    }

    #[instrument(skip(self), err)]
    pub fn search(
        &self,
        domain: String,
        query: String,
        format: QueryFormat,
        limit: Option<usize>,
    ) -> Result<String, Box<dyn Error>> {
        let di = self.domain_indexes.read();
        let mut out = String::new();

        if let Some((schema, _, reader, parser)) = di.get(&domain) {
            let query = parser.parse_query(&query)?;
            let searcher = reader.searcher();

            match format {
                QueryFormat::All => {
                    let docs = searcher.search(&*query, &tantivy::collector::DocSetCollector)?;

                    let mut lines = vec![];
                    tracing::info!(count = docs.len());

                    for doc_addr in docs {
                        let doc: TantivyDocument = searcher.doc(doc_addr)?;
                        let line = doc.to_json(&schema);
                        lines.push(line);
                    }

                    out = format!("[{}]", lines.join(","))
                }
                QueryFormat::Count => {
                    let count = searcher.search(&*query, &tantivy::collector::Count)?;

                    out = count.to_string();
                }
                QueryFormat::Top => {
                    if let Some(limit) = limit {
                        let docs = searcher
                            .search(&*query, &tantivy::collector::TopDocs::with_limit(limit))?;

                        let mut lines = vec![];
                        tracing::info!(count = docs.len());

                        for (_, doc_addr) in docs {
                            let doc: TantivyDocument = searcher.doc(doc_addr)?;
                            let line = doc.to_json(&schema);
                            lines.push(line);
                        }

                        out = format!("[{}]", lines.join(","))
                    }
                }
            }
        }

        Ok(out)
    }

    #[instrument(skip(dir, domain_indexes), err)]
    fn index(
        dir: &PathBuf,
        domain_indexes: &Arc<RwLock<BTreeMap<String, (Schema, Index, IndexReader, QueryParser)>>>,
    ) -> Result<(), Box<dyn Error>> {
        tracing::info!("indexing");

        let cursor_path = dir.join("cursor");

        let file_cursor_path = cursor_path.join("file");
        let file_cursor = match std::fs::read(&file_cursor_path) {
            Ok(file) => dir.join(std::str::from_utf8(&file)?),
            Err(_) => dir.join("."),
        };

        let line_cursor_path = cursor_path.join("line");
        let line_cursor = match std::fs::read(&line_cursor_path) {
            Ok(file) => std::str::from_utf8(&file)?.parse::<usize>()?,
            Err(_) => 0,
        };

        tracing::info!(file_cursor = %file_cursor.display(), line_cursor);

        let mut log_files = std::fs::read_dir(&dir)?
            .map(|res| res.map(|e| e.path()))
            .collect::<Result<Vec<_>, std::io::Error>>()?;
        log_files.sort();

        let di = domain_indexes.read();

        let mut writer_map = BTreeMap::new();

        for (domain, (schema, index, _, _)) in di.iter() {
            let writer: IndexWriter = index.writer(15_000_000)?;

            writer_map.insert(domain, (schema, writer));
        }

        let mut last_log_file = file_cursor.clone();
        let mut last_log_line = line_cursor;

        for log_file in log_files {
            if log_file < *file_cursor || log_file.is_dir() {
                continue;
            }

            tracing::info!(file = %log_file.display());

            last_log_file = log_file.clone();

            let file = File::open(&log_file)?;

            let reader = BufReader::new(file);

            for (i, line) in reader.lines().enumerate() {
                if log_file == *file_cursor {
                    if i < line_cursor {
                        continue;
                    }
                }

                let line = line?;
                let log = serde_json::from_str::<LogLine>(&line)?;

                let mut log_domain = String::new();

                if let Some(spans) = log.spans {
                    for span in spans {
                        if let Some(domain) = span.domain {
                            log_domain = domain;
                            break;
                        }
                    }
                }

                if let Some((schema, writer)) = writer_map.get(&log_domain) {
                    let doc = TantivyDocument::parse_json(&schema, &line)?;

                    writer.add_document(doc)?;
                }

                last_log_line = i;
            }

            tracing::info!(lines = last_log_line + 1);
        }

        tracing::info!(file_cursor = %last_log_file.display(), line_cursor = last_log_line);

        let mut file_cursor = File::create(&file_cursor_path)?;
        file_cursor.write_all(
            last_log_file
                .file_name()
                .unwrap()
                .to_str()
                .unwrap()
                .as_bytes(),
        )?;

        let mut line_cursor = File::create(&line_cursor_path)?;
        line_cursor.write_all(last_log_line.to_string().as_bytes())?;

        for (_, (_, mut writer)) in writer_map {
            writer.commit()?;
        }

        Ok(())
    }

    pub fn start(&self) -> Result<(JoinHandle<()>, JoinHandle<()>), Box<dyn Error>> {
        let cursor_path = self.dir.join("cursor");
        std::fs::create_dir_all(&cursor_path)?;

        let index_interval_s = self.index_interval_s;
        let cleanup_interval_s = self.cleanup_interval_s;
        let max_lines = self.max_lines;
        let max_ttl_days = self.max_ttl_days;
        let dir = self.dir.clone();
        let domain_indexes = self.domain_indexes.clone();

        // ?? a version of this could use notify to watch for new logs or log files
        // - https://docs.rs/notify/latest/notify/

        let index_handle = tokio::spawn(async move {
            loop {
                tokio::time::sleep(Duration::from_secs(index_interval_s)).await;

                match Self::index(&dir, &domain_indexes) {
                    Ok(()) => tracing::info!("indexed"),
                    Err(err) => tracing::error!(%err),
                };
            }
        });

        let cleanup_handle = tokio::spawn(async move {
            loop {
                tokio::time::sleep(Duration::from_secs(cleanup_interval_s)).await;
                tracing::info!(max_lines, max_ttl_days, "todo: cleanup");
            }
        });

        Ok((index_handle, cleanup_handle))
    }
}