h2kv 0.2.0

HTTP/2 interface for key-value storage systems
Documentation
use std::path::PathBuf;
use std::sync::{Arc, mpsc};

use anyhow::{Result, anyhow, bail};
use auto_args::AutoArgs;
use tokio::signal::unix::{SignalKind, signal as unix_signal};

#[derive(Debug, AutoArgs)]
struct Opt {
    /// print the package version and exit
    version: bool,
    /// directory to use for storage engine files
    storage_dir: Option<PathBuf>,
    /// listening port for TCP connections, default: 5928
    port: Option<i32>,
    /// directory to sync with the database on start and SIGHUP
    sync_dir: Option<PathBuf>,
    /// write to the synchronized directory on exit and SIGHUP
    sync_write: bool,
    /// fork into background process
    daemon: bool,
    /// PID file, ignored unless --daemon is set
    pidfile: Option<PathBuf>,
    /// file to send log messages, ignored unless --daemon is set
    log_filename: Option<PathBuf>,
}

impl TryFrom<Opt> for h2kv::Config {
    type Error = anyhow::Error;

    fn try_from(value: Opt) -> std::result::Result<Self, Self::Error> {
        if value.version {
            println!("{} {}", env!("CARGO_PKG_NAME"), env!("CARGO_PKG_VERSION"));
            std::process::exit(0);
        }

        match value.storage_dir {
            None => bail!("storage-dir is required"),
            Some(dir) if !dir.as_path().is_dir() => {
                bail!("storage-dir {dir:?} is not a directory")
            }
            _ => (),
        }

        match value.sync_dir {
            Some(sync_dir) if !sync_dir.as_path().is_dir() => {
                bail!("sync-dir {:?} is not a directory", &sync_dir);
            }
            _ => (),
        }

        if value.sync_write && value.sync_dir.is_none() {
            bail!("no sync-dir specified for sync-write");
        }

        if value.pidfile.is_some() && !value.daemon {
            log::warn!(
                "'--pidfile {:?}' ignored because '--daemon' is not set",
                value.pidfile.as_ref().unwrap()
            );
        }

        if value.log_filename.is_some() && !value.daemon {
            log::warn!(
                "'--log-filename {:?}' ignored because '--daemon' is not set",
                value.log_filename.as_ref().unwrap()
            );
        }

        let sync_ignore = h2kv::IgnoreFilter::try_from_env()?;
        if sync_ignore.is_active() {
            log::warn!("ignore filter {sync_ignore}");
        }

        Ok(Self {
            port: value.port.unwrap_or(5928),
            storage_dir: value.storage_dir.unwrap(),
            sync_dir: value.sync_dir,
            sync_write: value.sync_write,
            sync_ignore,
            daemon: value.daemon,
            pidfile: value.pidfile,
            log_filename: value.log_filename,
        })
    }
}

fn main() -> Result<()> {
    if cfg!(debug_assertions) {
        env_logger::Builder::from_default_env()
            .format_timestamp(None)
            .try_init()?;
    } else {
        env_logger::try_init()?;
    }

    help_intercept();
    let config: h2kv::Config = Opt::from_args().try_into()?;

    let (updates_tx, updates_rx) = mpsc::channel::<PathBuf>();

    let storage_dir = config.storage_dir.clone();
    let updates_tx_clone = updates_tx.clone();
    let lock_resources = move || -> Result<_, anyhow::Error> {
        let listener = std::net::TcpListener::bind(format!("127.0.0.1:{}", config.port))?;
        let db = h2kv::StorageFactory::try_create(&storage_dir, updates_tx_clone)?;
        Ok((listener, Arc::new(db)))
    };

    let (listener, db) = if config.daemon {
        match h2kv::runtime::spawn_daemon(&config, lock_resources)? {
            None => {
                log::trace!("daemon spawned. terminating parent");
                return Ok(());
            }
            Some(resources) => {
                log::trace!("daemon process started: {:?}", std::process::id());
                resources
            }
        }
    } else {
        lock_resources().map_err(|e| anyhow!("resource lock failure: {e}"))?
    };

    let files = h2kv::runtime::FilesystemActions {
        sync_dir: config.sync_dir.as_deref(),
        sync_write: config.sync_write,
        ignore: &config.sync_ignore,
        updates_rx: &updates_rx,
    };

    tokio::runtime::Builder::new_multi_thread()
        .enable_all()
        .build()
        .unwrap()
        .block_on(async {
            // port and async context are available, so reopen the socket in non-blocking mode
            let addr = listener.local_addr()?;
            drop(listener);
            let listener = tokio::net::TcpListener::bind(addr).await?;

            files.do_read(db.clone())?;

            loop {
                tokio::select! {
                    biased;
                    _ = signal(SignalKind::terminate()) => {
                        log::info!("received SIGTERM. exiting");
                        break;
                    },
                    _ = signal(SignalKind::interrupt()) => {
                        log::info!("received SIGINT. exiting");
                        break;
                    },
                    _ = signal(SignalKind::hangup()), if files.sync_dir.is_some() => {
                        log::info!(
                            "received SIGHUP. synchronizing db and filesystem ({:?})",
                            files.sync_dir.unwrap()
                        );
                        files.do_write(db.clone())?;
                        files.do_read(db.clone())?;
                    }
                    _ = h2kv::server::listen(&listener, db.clone()) => {},
                }
            }

            files.do_write(db)?;

            Ok(())
        })
}

async fn signal(kind: SignalKind) -> std::io::Result<()> {
    unix_signal(kind)?.recv().await;
    Ok(())
}

fn help_intercept() {
    let args = std::env::args().collect::<Vec<_>>();
    if args.contains(&"--help".to_string()) {
        let msg_lines = Opt::help();
        let mut msg_lines = msg_lines
            .lines()
            .filter(|l| *l != "For more information try --help")
            .collect::<Vec<_>>();
        let ignore_filter_description = format!(
            "{}: {}",
            h2kv::IgnoreFilter::ENV_NAME,
            h2kv::IgnoreFilter::ENV_DESCRIPTION
        );
        msg_lines.append(&mut vec![
            "Environment Variables:",
            &ignore_filter_description,
        ]);
        eprintln!("{}", msg_lines.join("\n"));

        std::process::exit(1);
    }
}