Allow reloading tilesets without server restart

Cargo.toml

@@ -34,6 +34,7 @@ rusqlite = "0.25"
 serde = { version = "1.0", features = ["derive"] }
 serde_json = "1.0"
 tokio = { version = "1.6", features = ["full"] }
+notify = "4.0.17"
 
 [dev-dependencies]
 tempdir = "0.3"

README.md

@@ -21,27 +21,22 @@ USAGE:
     mbtileserver [FLAGS] [OPTIONS]
 
 FLAGS:
-        --disable-preview    
-            Disable preview map
-    -h, --help               
-            Prints help information
-
-    -V, --version            
-            Prints version information
-
+        --allow-reload-api       Allow reloading tilesets with /reload endpoint
+        --allow-reload-signal    Allow reloading tilesets with a SIGHUP
+        --disable-preview        Disable preview map
+        --disable-watcher        Disable fs watcher for automatic tileset reload
+    -h, --help                   Prints help information
+    -V, --version                Prints version information
 
 OPTIONS:
-        --allowed-hosts <allowed_hosts>    
-            "*" matches all domains and ".<domain>" matches all subdomains for the given domain
-             [default: localhost, 127.0.0.1, [::1]]
-    -d, --directory <directory>            
-            Tiles directory
-             [default: ./tiles]
-    -H, --header <header>...               
-            Add custom header
-    -p, --port <port>                      
-            Server port
-             [default: 3000]
+        --allowed-hosts <allowed_hosts>        A comma-separated list of allowed hosts [default: localhost, 127.0.0.1,
+                                               [::1]]
+    -d, --directory <directory>                Tiles directory
+                                                [default: ./tiles]
+    -H, --header <header>...                   Add custom header
+    -p, --port <port>                          Server port
+                                                [default: 3000]
+        --reload-interval <reload_interval>    An interval at which tilesets get reloaded
 ```
 
 Run `mbtileserver` to start serving the mbtiles in a given folder. The default folder is `./tiles` and you can change it with `-d` flag.
@@ -53,6 +48,7 @@ You can adjust the log level by setting `RUST_LOG` environment variable. Possbil
 
 | Endpoint                                                     | Description                                                                    |
 |--------------------------------------------------------------|--------------------------------------------------------------------------------|
+| /reload                                                      | reloads tilesets from directory (if enabled with `--allow-reload`)             |
 | /services                                                    | lists all discovered and valid mbtiles in the tiles directory                  |
 | /services/\<path-to-tileset>                                 | shows tileset metadata                                                         |
 | /services/\<path-to-tileset>/map                             | tileset preview                                                                |

src/config.rs

@@ -1,19 +1,28 @@
-use std::collections::HashMap;
 use std::env;
 use std::path::PathBuf;
+use std::time::Duration;
 
 use clap::{crate_version, App, Arg, ArgMatches};
+use regex::Regex;
 
 use crate::errors::{Error, Result};
 use crate::tiles;
 
+lazy_static! {
+    static ref DURATION_RE: Regex = Regex::new(r"\d+[smhd]").unwrap();
+}
+
 #[derive(Clone, Debug)]
 pub struct Args {
-    pub tilesets: HashMap<String, tiles::TileMeta>,
+    pub tilesets: tiles::Tilesets,
     pub port: u16,
     pub allowed_hosts: Vec<String>,
     pub headers: Vec<(String, String)>,
     pub disable_preview: bool,
+    pub allow_reload_api: bool,
+    pub allow_reload_signal: bool,
+    pub reload_interval: Option<Duration>,
+    pub disable_watcher: bool,
 }
 
 pub fn get_app<'a, 'b>() -> App<'a, 'b> {
@@ -57,6 +66,28 @@ pub fn get_app<'a, 'b>() -> App<'a, 'b> {
                 .long("disable-preview")
                 .help("Disable preview map\n"),
         )
+        .arg(
+            Arg::with_name("allow_reload_api")
+                .long("allow-reload-api")
+                .help("Allow reloading tilesets with /reload endpoint\n"),
+        )
+        .arg(
+            Arg::with_name("allow_reload_signal")
+                .long("allow-reload-signal")
+                .help("Allow reloading tilesets with a SIGHUP\n"),
+        )
+        .arg(
+            Arg::with_name("reload_interval")
+                .long("reload-interval")
+                .help("An interval at which tilesets get reloaded")
+                .long_help("\"*\" in 1h30m format\n")
+                .takes_value(true),
+        )
+        .arg(
+            Arg::with_name("disable_watcher")
+                .long("disable-watcher")
+                .help("Disable fs watcher for automatic tileset reload\n"),
+        )
 }
 
 pub fn parse(matches: ArgMatches) -> Result<Args> {
@@ -107,14 +138,46 @@ pub fn parse(matches: ArgMatches) -> Result<Args> {
         }
     }
 
+    let reload_interval = match matches.value_of("reload_interval") {
+        Some(str) => {
+            let mut duration = Duration::ZERO;
+            for mat in DURATION_RE.find_iter(str) {
+                let mut mat = mat.as_str().to_owned();
+                let char = mat.chars().nth(mat.len() - 1).unwrap();
+                mat.truncate(mat.len() - 1);
+                let multiplier = match char {
+                    's' => 1,
+                    'm' => 60,
+                    'h' => 60 * 60,
+                    'd' => 60 * 60 * 24,
+                    _ => return Err(Error::Config("Invalid value for duration".to_string())),
+                };
+                let qty = match mat.parse::<u64>() {
+                    Ok(v) => v,
+                    Err(_) => return Err(Error::Config("Invalid value for duration".to_string())),
+                };
+                duration += Duration::from_secs(multiplier * qty);
+            }
+            Some(duration)
+        }
+        None => None,
+    };
+
     let disable_preview = matches.occurrences_of("disable_preview") != 0;
+    let allow_reload_api = matches.occurrences_of("allow_reload_api") != 0;
+    let allow_reload_signal = matches.occurrences_of("allow_reload_signal") != 0;
+    let disable_watcher = matches.occurrences_of("disable_watcher") != 0;
 
     Ok(Args {
         tilesets,
         port,
         allowed_hosts,
         headers,
         disable_preview,
+        allow_reload_api,
+        allow_reload_signal,
+        reload_interval,
+        disable_watcher,
     })
 }
 

src/main.rs

@@ -30,13 +30,7 @@ fn main() {
         }
     };
 
-    if let Err(e) = server::run(
-        args.port,
-        args.allowed_hosts,
-        args.headers,
-        args.disable_preview,
-        args.tilesets,
-    ) {
+    if let Err(e) = server::run(args) {
         error!("Server error: {}", e);
         std::process::exit(1);
     }

src/server.rs

@@ -1,37 +1,85 @@
+use crate::config::Args;
 use hyper::service::{make_service_fn, service_fn};
 use hyper::Server;
-use std::collections::HashMap;
+use notify::Watcher;
+use std::time::Duration;
 
 use crate::service;
-use crate::tiles::TileMeta;
 
 #[tokio::main]
-pub async fn run(
-    port: u16,
-    allowed_hosts: Vec<String>,
-    headers: Vec<(String, String)>,
-    disable_preview: bool,
-    tilesets: HashMap<String, TileMeta>,
-) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
-    let addr = ([0, 0, 0, 0], port).into();
+pub async fn run(args: Args) -> Result<(), Box<dyn std::error::Error + Send + Sync>> {
+    let addr = ([0, 0, 0, 0], args.port).into();
     let server = Server::try_bind(&addr)?;
 
-    let service = make_service_fn(move |_conn| {
-        let tilesets = tilesets.clone();
-        let allowed_hosts = allowed_hosts.clone();
-        let headers = headers.clone();
-        async move {
-            Ok::<_, hyper::Error>(service_fn(move |req| {
-                service::get_service(
-                    req,
-                    tilesets.clone(),
-                    allowed_hosts.clone(),
-                    headers.clone(),
-                    disable_preview,
-                )
-            }))
-        }
-    });
+    let service = {
+        let args = args.clone();
+        make_service_fn(move |_conn| {
+            let tilesets = args.tilesets.clone();
+            let allowed_hosts = args.allowed_hosts.clone();
+            let headers = args.headers.clone();
+            let disable_preview = args.disable_preview;
+            let allow_reload_api = args.allow_reload_api;
+            async move {
+                Ok::<_, hyper::Error>(service_fn(move |req| {
+                    service::get_service(
+                        req,
+                        tilesets.clone(),
+                        allowed_hosts.clone(),
+                        headers.clone(),
+                        disable_preview,
+                        allow_reload_api,
+                    )
+                }))
+            }
+        })
+    };
+
+    if args.allow_reload_signal {
+        let tilesets = args.tilesets.clone();
+        println!("Reloading on SIGHUP");
+        tokio::spawn(async move {
+            let mut handler =
+                tokio::signal::unix::signal(tokio::signal::unix::SignalKind::hangup()).unwrap();
+            loop {
+                handler.recv().await;
+                println!("Caught SIGHUP, reloading tilesets");
+                tilesets.reload();
+            }
+        });
+    }
+
+    if let Some(interval) = args.reload_interval {
+        let tilesets = args.tilesets.clone();
+        println!("Reloading every {} seconds", interval.as_secs());
+        tokio::spawn(async move {
+            let mut timer = tokio::time::interval(interval);
+            loop {
+                timer.tick().await;
+                tilesets.reload();
+            }
+        });
+    }
+
+    if !args.disable_watcher {
+        let tilesets = args.tilesets.clone();
+        println!("Watching FS for changes on {:?}", tilesets.get_path());
+        drop(std::thread::spawn(move || {
+            let (tx, rx) = std::sync::mpsc::channel();
+            let mut watcher = notify::watcher(tx, Duration::from_secs(10)).unwrap();
+            watcher
+                .watch(&args.tilesets.get_path(), notify::RecursiveMode::Recursive)
+                .unwrap();
+            loop {
+                match rx.recv() {
+                    Ok(_) => {
+                        println!("Tileset directory changed, reloading");
+                        tilesets.reload()
+                    }
+                    Err(e) => println!("watch error: {:?}", e),
+                }
+            }
+        }));
+    }
 
     println!("Listening on http://{}", addr);
     server.serve(service).await?;