1 files changed, 73 insertions, 64 deletions
diff --git a/server/replaytool/src/main.rs b/server/replaytool/src/main.rs
index e8dae5ef..f92ac90a 100644
--- a/server/replaytool/src/main.rs
+++ b/server/replaytool/src/main.rs
@@ -23,13 +23,14 @@ use hurrycurry_protocol::{PacketC, PacketS};
 use log::{debug, info, warn, LevelFilter};
 use serde::{Deserialize, Serialize};
 use std::{
-    path::PathBuf,
-    time::{Instant, SystemTime},
+    path::{Path, PathBuf},
+    time::{Duration, Instant, SystemTime},
 };
 use tokio::{
     fs::File,
     io::{AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter},
     net::TcpListener,
+    time::sleep,
 };
 use tokio_tungstenite::tungstenite::Message;
 
@@ -68,69 +69,21 @@ async fn main() -> anyhow::Result<()> {
             output,
             r#loop,
         } => loop {
-            let mut file = BufWriter::new(ZstdEncoder::new(BufWriter::new(
-                File::create(&if r#loop {
-                    output.join(format!(
-                        "replay-{}",
-                        SystemTime::now()
-                            .duration_since(SystemTime::UNIX_EPOCH)
-                            .unwrap()
-                            .as_secs()
-                    ))
-                } else {
-                    output.clone()
-                })
-                .await?,
-            )));
-            info!("connecting to {url:?}...");
-            let (mut sock, _) = tokio_tungstenite::connect_async(&url).await?;
-            info!("starting recording.");
-            let start = Instant::now();
-
-            while let Some(Ok(message)) = sock.next().await {
-                match message {
-                    Message::Text(line) => {
-                        let packet: PacketC = match serde_json::from_str(&line) {
-                            Ok(p) => p,
-                            Err(e) => {
-                                warn!("invalid packet: {e}");
-                                break;
-                            }
-                        };
-                        debug!("<- {packet:?}");
-
-                        let is_end = matches!(packet, PacketC::SetIngame { state: false, .. });
-
-                        file.write_all(
-                            format!(
-                                "{}\n",
-                                serde_json::to_string(&Event {
-                                    ts: start.elapsed().as_secs_f64(),
-                                    packet: packet
-                                })
-                                .unwrap()
-                            )
-                            .as_bytes(),
-                        )
-                        .await?;
-
-                        if is_end {
-                            info!("stopping replay...");
-                            break;
-                        }
-                    }
-                    Message::Close(_) => break,
-                    _ => (),
-                }
+            let out = if r#loop {
+                output.join(format!(
+                    "replay-{}",
+                    SystemTime::now()
+                        .duration_since(SystemTime::UNIX_EPOCH)
+                        .unwrap()
+                        .as_secs()
+                ))
+            } else {
+                output.clone()
+            };
+            if let Err(e) = do_record(&out, &url).await {
+                warn!("recording failed: {e}");
+                sleep(Duration::from_secs(1)).await;
             }
-            file.flush().await?;
-            let mut file = file.into_inner();
-            file.flush().await?;
-            let mut file = file.into_inner();
-            file.flush().await?;
-            let mut file = file.into_inner();
-            file.flush().await?;
-            info!("done");
             if r#loop {
                 info!("restarting...");
             } else {
@@ -213,3 +166,59 @@ async fn main() -> anyhow::Result<()> {
     }
     Ok(())
 }
+
+pub async fn do_record(output: &Path, url: &str) -> anyhow::Result<()> {
+    let mut file = BufWriter::new(ZstdEncoder::new(BufWriter::new(
+        File::create(&output).await?,
+    )));
+    info!("connecting to {url:?}...");
+    let (mut sock, _) = tokio_tungstenite::connect_async(url).await?;
+    info!("starting recording.");
+    let start = Instant::now();
+
+    while let Some(Ok(message)) = sock.next().await {
+        match message {
+            Message::Text(line) => {
+                let packet: PacketC = match serde_json::from_str(&line) {
+                    Ok(p) => p,
+                    Err(e) => {
+                        warn!("invalid packet: {e}");
+                        break;
+                    }
+                };
+                debug!("<- {packet:?}");
+
+                let is_end = matches!(packet, PacketC::SetIngame { state: false, .. });
+
+                file.write_all(
+                    format!(
+                        "{}\n",
+                        serde_json::to_string(&Event {
+                            ts: start.elapsed().as_secs_f64(),
+                            packet: packet
+                        })
+                        .unwrap()
+                    )
+                    .as_bytes(),
+                )
+                .await?;
+
+                if is_end {
+                    info!("stopping replay...");
+                    break;
+                }
+            }
+            Message::Close(_) => break,
+            _ => (),
+        }
+    }
+    file.flush().await?;
+    let mut file = file.into_inner();
+    file.flush().await?;
+    let mut file = file.into_inner();
+    file.flush().await?;
+    let mut file = file.into_inner();
+    file.flush().await?;
+    info!("done");
+    Ok(())
+}