Skip to content

Commit

Permalink
リファクタリング
Browse files Browse the repository at this point in the history
  • Loading branch information
mousecrusher2 committed Apr 17, 2024
1 parent 6624e93 commit abfd946
Show file tree
Hide file tree
Showing 2 changed files with 67 additions and 56 deletions.
99 changes: 43 additions & 56 deletions src/main.rs
Original file line number Diff line number Diff line change
@@ -1,14 +1,16 @@
mod backup;
mod util;

use clap::Parser;
use regex::Regex;
use std::path::PathBuf;
use tokio::{
fs,
io::{self, AsyncBufReadExt, AsyncWriteExt, BufReader, BufWriter},
io::{self, AsyncBufReadExt, BufReader},
process::Command,
time,
time::{self, sleep},
};
use util::Writer;

const COMMAND_PREFIX: &str = "!";

Expand All @@ -33,6 +35,17 @@ struct Cli {
#[tokio::main(flavor = "current_thread")]
async fn main() {
let cli = Cli::parse();
if !cli.backup_folder.exists() {
fs::create_dir_all(&cli.backup_folder)
.await
.expect("Failed to create backup folder");
}

if !cli.backup_folder.is_dir() {
eprintln!("Backup folder must be a directory");
std::process::exit(1);
}

let mut child = Command::new(&cli.server_path)
.env("LD_LIBRARY_PATH", cli.server_path.parent().unwrap())
.args(cli.server_options)
Expand All @@ -46,25 +59,18 @@ async fn main() {
// If the server is too slow to start, fail to read from stdout
time::sleep(time::Duration::from_secs(3)).await;

if !cli.backup_folder.exists() {
fs::create_dir_all(&cli.backup_folder)
.await
.expect("Failed to create backup folder");
}

if !cli.backup_folder.is_dir() {
eprintln!("Backup folder must be a directory");
std::process::exit(1);
}

let mut self_stdin = BufReader::new(io::stdin()).lines();
let mut self_stdout = BufWriter::new(io::stdout());
let mut self_stdout = Writer::new(io::stdout());
let mut server_stdout = BufReader::new(child.stdout.take().unwrap()).lines();
let mut server_stdin = child.stdin.take().unwrap();
let mut server_stdin = Writer::new(child.stdin.take().unwrap());

let backuper = backup::Backuper::new(
cli.backup_folder,
cli.server_path.parent().unwrap().to_path_buf().join("worlds"),
cli.server_path
.parent()
.unwrap()
.to_path_buf()
.join("worlds"),
cli.backup_count.unwrap_or(1),
);

Expand All @@ -91,9 +97,7 @@ async fn main() {

line = server_stdout.next_line() => {
if let Some(line) = line.expect("Failed to read from server stdout") {
self_stdout.write_all(line.as_bytes()).await.expect("Failed to write to stdout");
self_stdout.write_u8(b'\n').await.expect("Failed to write to stdout");
self_stdout.flush().await.unwrap();
self_stdout.writeln_flush(line.as_bytes()).await.expect("Failed to write to stdout");
}
}

Expand All @@ -103,11 +107,8 @@ async fn main() {
.expect("stdin closed. Don't you redirect stdin?");
if line.starts_with(COMMAND_PREFIX) {
if backup_command_re.is_match(line.strip_prefix(COMMAND_PREFIX).unwrap()) {
self_stdout.write_all(b"Backup started\n").await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
server_stdin.write_all(b"save hold\n").await.expect("Failed to write to server stdin");
server_stdin.write_all(b"save query\n").await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
self_stdout.write_flush(b"Backup started\n").await.expect("Failed to write to stdout");
server_stdin.write_flush(b"save hold\nsave query\n").await.expect("Failed to write to server stdin");

server_stdout.next_line().await.expect("Failed to read from server stdout");
loop {
Expand All @@ -117,44 +118,32 @@ async fn main() {
if ready_backups_re.is_match(&line) {
break;
}
self_stdout.write_all(line.as_bytes()).await.expect("Failed to write to stdout");
self_stdout.write_u8(b'\n').await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
server_stdin.write_all(b"save query\n").await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
tokio::time::sleep(time::Duration::from_millis(300)).await;
self_stdout.writeln_flush(line.as_bytes()).await.expect("Failed to write to stdout");
server_stdin.write_flush(b"save query\n").await.expect("Failed to write to server stdin");
sleep(time::Duration::from_millis(300)).await;
}
let line = server_stdout.next_line().await
.expect("Failed to read from server stdout")
.expect("Server stdout closed");
backuper.backup(line).await;
server_stdin.write_all(b"save resume\n").await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
server_stdin.write_flush(b"save resume\n").await.expect("Failed to write to server stdin");

self_stdout.write_all(b"Backup finished\n").await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
self_stdout.write_flush(b"Backup finished\n").await.expect("Failed to write to stdout");

} else {
self_stdout.write_all(b"Unknown command\n").await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
self_stdout.write_flush(b"Unknown command\n").await.expect("Failed to write to stdout");
}
} else if command_re.is_match(&line) {
// save command is not allowed
self_stdout.write_all(b"Don't use \"save xxxx\" command").await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
self_stdout.write_flush(b"Don't use \"save xxxx\" command\n").await.expect("Failed to write to stdout");
} else {
server_stdin.write_all(line.as_bytes()).await.expect("Failed to write to server stdin");
server_stdin.write_u8(b'\n').await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
server_stdin.writeln_flush(line.as_bytes()).await.expect("Failed to write to server stdin");
}
}

_ = backup_interval.tick() => {
self_stdout.write_all(b"Backup started\n").await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
server_stdin.write_all(b"save hold\n").await.expect("Failed to write to server stdin");
server_stdin.write_all(b"save query\n").await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
self_stdout.write_flush(b"Backup started\n").await.expect("Failed to write to stdout");
server_stdin.write_flush(b"save hold\nsave query\n").await.expect("Failed to write to server stdin");

server_stdout.next_line().await.expect("Failed to read from server stdout");

Expand All @@ -165,22 +154,17 @@ async fn main() {
if ready_backups_re.is_match(&line) {
break;
}
self_stdout.write_all(line.as_bytes()).await.expect("Failed to write to stdout");
self_stdout.write_u8(b'\n').await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
self_stdout.writeln_flush(line.as_bytes()).await.expect("Failed to write to stdout");

server_stdin.write_all(b"save query\n").await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
server_stdin.write_flush(b"save query\n").await.expect("Failed to write to server stdin");
}
let line = server_stdout.next_line().await
.expect("Failed to read from server stdout")
.expect("Server stdout closed");
backuper.backup(line).await;
server_stdin.write_all(b"save resume\n").await.expect("Failed to write to server stdin");
server_stdin.flush().await.expect("Failed to flush server stdin");
server_stdin.write_flush(b"save resume\n").await.expect("Failed to write to server stdin");

self_stdout.write_all(b"Backup finished\n").await.expect("Failed to write to stdout");
self_stdout.flush().await.expect("Failed to flush stdout");
self_stdout.write_flush(b"Backup finished\n").await.expect("Failed to write to stdout");
}

}
Expand All @@ -195,6 +179,9 @@ mod test {
fn format_chrono() {
let now = Local::now();
let formatted = now.format("%F %H_%M_%S%.f %z").to_string();
assert_eq!(now, DateTime::parse_from_str(&formatted, "%F %H_%M_%S%.f %z").unwrap());
assert_eq!(
now,
DateTime::parse_from_str(&formatted, "%F %H_%M_%S%.f %z").unwrap()
);
}
}
24 changes: 24 additions & 0 deletions src/util.rs
Original file line number Diff line number Diff line change
@@ -0,0 +1,24 @@
use tokio::io::{self, AsyncWriteExt};

pub(crate) struct Writer<T: AsyncWriteExt + Unpin> {
writer: T,
}

impl<T: AsyncWriteExt + Unpin> Writer<T> {
pub(crate) fn new(writer: T) -> Self {
Self { writer }
}

pub(crate) async fn write_flush(&mut self, data: &[u8]) -> io::Result<()> {
self.writer.write_all(data).await?;
self.writer.flush().await?;
Ok(())
}

pub(crate) async fn writeln_flush(&mut self, data: &[u8]) -> io::Result<()> {
self.writer.write_all(data).await?;
self.writer.write_u8(b'\n').await?;
self.writer.flush().await?;
Ok(())
}
}

0 comments on commit abfd946

Please sign in to comment.