~neon/activity-graph

ac0de63c7532de1a91dd00e210f2a0bf1275d704 — Jens Pitkanen 2 years ago 3acc6f5
Add --cache-file
2 files changed, 112 insertions(+), 10 deletions(-)

M src/main.rs
M src/server.rs
M src/main.rs => src/main.rs +13 -6
@@ 143,6 143,12 @@ enum CommandArgs {
        /// html and css
        #[structopt(long, default_value = "1")]
        cache_lifetime: u64,
        /// A file that will be used as backup storage for the cache
        /// (useful when you want to keep serving the previous cached
        /// version after restarting the server, to avoid a period of
        /// unresponsiveness)
        #[structopt(long)]
        cache_file: Option<PathBuf>,
    },
}



@@ 150,7 156,7 @@ fn main() {
    let start_time = time::Instant::now();
    let args = Args::from_args();

    if let Some(command) = &args.command {
    if let Some(command) = args.command {
        match command {
            CommandArgs::Generate {
                verbosity,


@@ 159,7 165,7 @@ fn main() {
                html,
                css,
            } => {
                log::set_verbosity(verbosity);
                log::set_verbosity(&verbosity);

                let write_to_file = |path: &Path, s: String, name: &str| {
                    let mut writer = File::create(path).map(BufWriter::new);


@@ 193,8 199,8 @@ fn main() {
            }

            CommandArgs::Stdout { verbosity, gen } => {
                log::set_verbosity(verbosity);
                println!("{}", render::ascii(&generate_years(gen)));
                log::set_verbosity(&verbosity);
                println!("{}", render::ascii(&generate_years(&gen)));
            }

            #[cfg(feature = "server")]


@@ 204,9 210,10 @@ fn main() {
                ext,
                host,
                cache_lifetime,
                cache_file,
            } => {
                log::set_verbosity(verbosity);
                server::run(&gen, &ext, *host, *cache_lifetime);
                log::set_verbosity(&verbosity);
                server::run(&gen, &ext, cache_file, host, cache_lifetime);
            }
        }
    }

M src/server.rs => src/server.rs +99 -4
@@ 5,6 5,8 @@ use tokio::runtime::Runtime;
use tokio::task;

use std::convert::Infallible;
use std::fs::File;
use std::io::{self, BufReader, BufWriter, Read, Write};
use std::net::SocketAddr;
use std::path::PathBuf;
use std::sync::atomic::{AtomicBool, Ordering};


@@ 20,6 22,9 @@ lazy_static::lazy_static! {
    static ref LAST_CACHE: RwLock<Instant> = RwLock::new(Instant::now());
    static ref CACHE_LIFETIME: RwLock<Duration> = RwLock::new(Duration::from_secs(0));

    // A backup of the current CACHED_HTML and CACHED_CSS values on
    // disk. Encoded in the order: <html> <CACHE_FILE_SPLITTER> <css>
    static ref CACHE_FILE: RwLock<Option<PathBuf>> = RwLock::new(None);
    static ref CACHED_HTML: RwLock<String> = RwLock::new(String::new());
    static ref CACHED_CSS: RwLock<String> = RwLock::new(String::new());
}


@@ 29,17 34,29 @@ static CACHE_INITIALIZED: AtomicBool = AtomicBool::new(false);

static INDEX_PATHS: &[&str] = &["/", "/index.html", "/index.htm", ""];

pub fn run(gen: &GenerationData, ext: &ExternalResources, host: SocketAddr, cache_lifetime: u64) {
// This is invalid UTF-8, and so can be used as a delimiter between
// Strings, as Strings are always valid UTF-8.
const CACHE_FILE_SPLITTER: u8 = 0xFE;

pub fn run(
    gen: &GenerationData,
    ext: &ExternalResources,
    cache_file: Option<PathBuf>,
    host: SocketAddr,
    cache_lifetime: u64,
) {
    log::verbose_println(&format!("starting server on {}...", host), true);

    if let (Ok(mut gen_), Ok(mut ext_), Ok(mut lifetime), Ok(mut last_cache)) = (
    if let (Ok(mut gen_), Ok(mut ext_), Ok(mut cache_file_), Ok(mut lifetime), Ok(mut last_cache)) = (
        GENERATION_DATA.write(),
        EXTERNAL_HTML.write(),
        CACHE_FILE.write(),
        CACHE_LIFETIME.write(),
        LAST_CACHE.write(),
    ) {
        *gen_ = gen.clone();
        *ext_ = ext.clone();
        *cache_file_ = cache_file;
        *lifetime = Duration::from_secs(cache_lifetime);
        *last_cache = Instant::now() - Duration::from_secs(cache_lifetime * 2);
    } else {


@@ 106,6 123,22 @@ async fn refresh_caches() {
            && !REFRESHING_CACHE.compare_and_swap(false, true, Ordering::Relaxed)
        {
            log::verbose_println("refreshing cache...", false);

            // Load from cache file if the cache has not been
            // initialized yet (if it exists)
            if !CACHE_INITIALIZED.load(Ordering::Relaxed) {
                if let Some((html, css)) = read_cache_file() {
                    if let (Ok(mut html_cache), Ok(mut css_cache)) =
                        (CACHED_HTML.write(), CACHED_CSS.write())
                    {
                        *html_cache = html;
                        *css_cache = css;
                        CACHE_INITIALIZED.store(true, Ordering::Relaxed);
                        log::println("initialized cache from cache file");
                    }
                }
            }

            let start = Instant::now();
            if let (Ok(gen), Ok(ext)) = (GENERATION_DATA.read(), EXTERNAL_HTML.read()) {
                let years = generate_years(&gen);


@@ 113,6 146,17 @@ async fn refresh_caches() {
                let css_path = PathBuf::from("/activity-graph.css");
                let output_html = render::html(&ext, &html_path, Some(&css_path), &years);
                let output_css = render::css(&ext);

                let (cache_html, cache_css) = (output_html.clone(), output_css.clone());
                task::spawn(async move {
                    if let Err(err) = write_cache_file(&cache_html, &cache_css) {
                        log::println(&format!(
                            "error: ran into an IO error while writing cache file: {}",
                            err
                        ));
                    }
                });

                if let Ok(mut html) = CACHED_HTML.write() {
                    *html = output_html;
                }


@@ 124,8 168,9 @@ async fn refresh_caches() {
                }
            }
            log::println(&format!("updated cache, took {:?}", Instant::now() - start));
            REFRESHING_CACHE.store(false, Ordering::Relaxed);
            CACHE_INITIALIZED.store(true, Ordering::Relaxed);

            REFRESHING_CACHE.store(false, Ordering::Relaxed); // Allow future refreshes
            CACHE_INITIALIZED.store(true, Ordering::Relaxed); // Allow early requests to complete
        }
    });



@@ 134,3 179,53 @@ async fn refresh_caches() {
        task::yield_now().await;
    }
}

fn write_cache_file(html: &str, css: &str) -> Result<(), io::Error> {
    if let Ok(cache_file) = CACHE_FILE.read() {
        if cache_file.is_some() {
            log::verbose_println("writing cache file...", true);
            let file = File::create(cache_file.as_ref().unwrap())?;
            let mut writer = BufWriter::new(file);
            write!(writer, "ACTIVITY-GRAPH-CACHE-FILE")?;
            writer.write(&[CACHE_FILE_SPLITTER])?;
            write!(writer, "{}", html)?;
            writer.write(&[CACHE_FILE_SPLITTER])?;
            write!(writer, "{}", css)?;
            drop(writer); // This should flush out the file write
            log::verbose_println("wrote cache file", false);
        }
    }
    Ok(())
}

fn read_cache_file() -> Option<(String, String)> {
    if let Ok(cache_file) = CACHE_FILE.read() {
        if cache_file.is_some() {
            let file = File::open(cache_file.as_ref().unwrap());
            match file {
                Ok(file) => {
                    let mut reader = BufReader::new(file);
                    let mut bytes = Vec::new();
                    if reader.read_to_end(&mut bytes).is_ok() {
                        // Split at CACHE_FILE_SPLITTER and return the
                        // parts between as `&str`s.
                        let parts: Vec<&str> = bytes
                            .split(|b| *b == CACHE_FILE_SPLITTER)
                            .filter_map(|bytes: &[u8]| std::str::from_utf8(bytes).ok())
                            .collect();
                        if parts.len() == 3 {
                            let (magic, html, css) = (parts[0], parts[1], parts[2]);
                            if magic == "ACTIVITY-GRAPH-CACHE-FILE" {
                                return Some((html.to_string(), css.to_string()));
                            }
                        }
                    }
                }
                Err(err) => {
                    log::println(&format!("error: could not read cache file: {}", err));
                }
            }
        }
    }
    None
}