1 files changed, 27 insertions, 6 deletions
diff --git a/src/main.rs b/src/main.rs
index e4a17e8..0441cc6 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -4,13 +4,15 @@ pub mod webui;
 pub mod webui_ws;
 pub mod worker_ws;
 
-use anyhow::Result;
+use anyhow::{Result, anyhow};
 use api::{api_complete_json, api_loading_json, api_queue_json};
 use axum::{Router, routing::get};
-use log::{debug, info};
+use log::{debug, error, info};
 use serde_json::{Map, Value};
 use std::{
     collections::{HashMap, HashSet},
+    env::args,
+    path::Path,
     sync::Arc,
     time::Instant,
 };
@@ -40,6 +42,8 @@ pub struct State {
     workers: HashMap<WorkerID, Worker>,
     webui_broadcast: broadcast::Sender<Arc<WebuiEvent>>,
 
+    config: Value,
+
     metadata: HashMap<String, Map<String, Value>>,
     queue: HashSet<String>,
     loading: HashSet<String>,
@@ -50,6 +54,7 @@ pub struct State {
 async fn main() -> Result<()> {
     env_logger::init_from_env("LOG");
     let mut state = State::default();
+    state.load_config().await?;
     state.load().await?;
     let router = Router::new()
         .route("/", get(webui))
@@ -69,6 +74,7 @@ async fn main() -> Result<()> {
 impl Default for State {
     fn default() -> Self {
         Self {
+            config: Default::default(),
             worker_id_counter: Default::default(),
             workers: Default::default(),
             webui_broadcast: broadcast::channel(1024).0,
@@ -81,13 +87,28 @@ impl Default for State {
 }
 
 impl State {
+    pub async fn load_config(&mut self) -> Result<()> {
+        let path = args()
+            .nth(1)
+            .ok_or(anyhow!("first argument is config path"))?;
+
+        self.config = serde_yml::from_str::<serde_json::Value>(&read_to_string(path).await?)?;
+        Ok(())
+    }
     pub async fn load(&mut self) -> Result<()> {
         debug!("loading state");
         let t = Instant::now();
-        self.metadata = serde_json::from_str(&read_to_string("metadata.json").await?)?;
-        self.queue = serde_json::from_str(&read_to_string("queue.json").await?)?;
-        self.complete = serde_json::from_str(&read_to_string("complete.json").await?)?;
-        info!("state loaded (took {:?})", t.elapsed());
+        if AsRef::<Path>::as_ref("metadata.json").exists()
+            && AsRef::<Path>::as_ref("queue.json").exists()
+            && AsRef::<Path>::as_ref("complete.json").exists()
+        {
+            self.metadata = serde_json::from_str(&read_to_string("metadata.json").await?)?;
+            self.queue = serde_json::from_str(&read_to_string("queue.json").await?)?;
+            self.complete = serde_json::from_str(&read_to_string("complete.json").await?)?;
+            info!("state loaded (took {:?})", t.elapsed());
+        } else {
+            error!("some state files are missing, skipping load")
+        }
         Ok(())
     }
     pub async fn save(&mut self) -> Result<()> {