parent
35f1de6c5e
commit
ecc8854e44
GPG Key ID:
284488987C31F630
@ -1,2 +1,3 @@
|
|||||||
/target
|
/target
|
||||||
*~
|
*~
|
||||||
|
chain.dat
|
||||||
|
|||||||
@ -0,0 +1,6 @@
|
|||||||
|
bindpoint = '127.0.0.1:8001'
|
||||||
|
file = 'chain.dat'
|
||||||
|
max_content_length = 4194304
|
||||||
|
max_gen_size = 256
|
||||||
|
#save_interval_secs = 2
|
||||||
|
trust_x_forwarded_for = false
|
||||||
@ -0,0 +1,108 @@
|
|||||||
|
//! Server config
|
||||||
|
use super::*;
|
||||||
|
use std::{
|
||||||
|
net::SocketAddr,
|
||||||
|
path::Path,
|
||||||
|
io,
|
||||||
|
borrow::Cow,
|
||||||
|
num::NonZeroU64,
|
||||||
|
};
|
||||||
|
use tokio::{
|
||||||
|
fs::OpenOptions,
|
||||||
|
prelude::*,
|
||||||
|
time::Duration,
|
||||||
|
io::BufReader,
|
||||||
|
};
|
||||||
|
|
||||||
|
pub const DEFAULT_FILE_LOCATION: &'static str = "markov.toml";
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Hash, Serialize, Deserialize)]
|
||||||
|
pub struct Config
|
||||||
|
{
|
||||||
|
pub bindpoint: SocketAddr,
|
||||||
|
pub file: String,
|
||||||
|
pub max_content_length: u64,
|
||||||
|
pub max_gen_size: usize,
|
||||||
|
pub save_interval_secs: Option<NonZeroU64>,
|
||||||
|
pub trust_x_forwarded_for: bool,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Default for Config
|
||||||
|
{
|
||||||
|
#[inline]
|
||||||
|
fn default() -> Self
|
||||||
|
{
|
||||||
|
Self {
|
||||||
|
bindpoint: ([127,0,0,1], 8001).into(),
|
||||||
|
file: "chain.dat".to_owned(),
|
||||||
|
max_content_length: 1024 * 1024 * 4,
|
||||||
|
max_gen_size: 256,
|
||||||
|
save_interval_secs: Some(unsafe{NonZeroU64::new_unchecked(2)}),
|
||||||
|
trust_x_forwarded_for: false,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
impl Config
|
||||||
|
{
|
||||||
|
pub fn save_interval(&self) -> Option<Duration>
|
||||||
|
{
|
||||||
|
self.save_interval_secs.map(|x| Duration::from_secs(x.into()))
|
||||||
|
}
|
||||||
|
pub async fn load(from: impl AsRef<Path>) -> io::Result<Self>
|
||||||
|
{
|
||||||
|
let file = OpenOptions::new()
|
||||||
|
.read(true)
|
||||||
|
.open(from).await?;
|
||||||
|
|
||||||
|
let mut buffer= String::new();
|
||||||
|
let reader = BufReader::new(file);
|
||||||
|
let mut lines = reader.lines();
|
||||||
|
while let Some(line) = lines.next_line().await? {
|
||||||
|
buffer.push_str(&line[..]);
|
||||||
|
buffer.push('\n');
|
||||||
|
}
|
||||||
|
toml::de::from_str(&buffer[..]).map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn save(&self, to: impl AsRef<Path>) -> io::Result<()>
|
||||||
|
{
|
||||||
|
let config = toml::ser::to_string_pretty(self).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||||||
|
let mut file = OpenOptions::new()
|
||||||
|
.write(true)
|
||||||
|
.create(true)
|
||||||
|
.truncate(true)
|
||||||
|
.open(to).await?;
|
||||||
|
file.write_all(config.as_bytes()).await?;
|
||||||
|
file.shutdown().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Try to load config file specified by args, or default config file
|
||||||
|
pub fn load() -> impl futures::future::Future<Output =Option<Config>>
|
||||||
|
{
|
||||||
|
load_args(std::env::args().skip(1))
|
||||||
|
}
|
||||||
|
|
||||||
|
async fn load_args<I: Iterator<Item=String>>(mut from: I) -> Option<Config>
|
||||||
|
{
|
||||||
|
let place = if let Some(arg) = from.next() {
|
||||||
|
trace!("File {:?} provided", arg);
|
||||||
|
Cow::Owned(arg)
|
||||||
|
} else {
|
||||||
|
warn!("No config file provided. Using default location {:?}", DEFAULT_FILE_LOCATION);
|
||||||
|
Cow::Borrowed(DEFAULT_FILE_LOCATION)
|
||||||
|
};
|
||||||
|
|
||||||
|
match Config::load(place.as_ref()).await {
|
||||||
|
Ok(cfg) => {
|
||||||
|
info!("Loaded config file {:?}", place);
|
||||||
|
Some(cfg)
|
||||||
|
},
|
||||||
|
Err(err) => {
|
||||||
|
error!("Failed to load config file from {:?}: {}", place, err);
|
||||||
|
None
|
||||||
|
},
|
||||||
|
}
|
||||||
|
}
|
||||||
@ -0,0 +1,89 @@
|
|||||||
|
//! Saving and loading chain
|
||||||
|
use super::*;
|
||||||
|
use std::{
|
||||||
|
sync::Arc,
|
||||||
|
path::{
|
||||||
|
Path,
|
||||||
|
},
|
||||||
|
io,
|
||||||
|
};
|
||||||
|
use tokio::{
|
||||||
|
time::{
|
||||||
|
self,
|
||||||
|
Duration,
|
||||||
|
},
|
||||||
|
fs::{
|
||||||
|
OpenOptions,
|
||||||
|
},
|
||||||
|
prelude::*,
|
||||||
|
};
|
||||||
|
use futures::{
|
||||||
|
future::{
|
||||||
|
OptionFuture,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
use lzzzz::{
|
||||||
|
lz4f::{
|
||||||
|
self,
|
||||||
|
AsyncWriteCompressor,
|
||||||
|
PreferencesBuilder,
|
||||||
|
AsyncReadDecompressor,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
|
||||||
|
const SAVE_INTERVAL: Option<Duration> = Some(Duration::from_secs(2));
|
||||||
|
|
||||||
|
pub async fn save_now(chain: &Chain<String>, to: impl AsRef<Path>) -> io::Result<()>
|
||||||
|
{
|
||||||
|
debug!("Saving chain to {:?}", to.as_ref());
|
||||||
|
let file = OpenOptions::new()
|
||||||
|
.write(true)
|
||||||
|
.create(true)
|
||||||
|
.truncate(true)
|
||||||
|
.open(to).await?;
|
||||||
|
let chain = serde_cbor::to_vec(chain).map_err(|e| io::Error::new(io::ErrorKind::InvalidData, e))?;
|
||||||
|
|
||||||
|
let mut file = AsyncWriteCompressor::new(file, PreferencesBuilder::new()
|
||||||
|
.compression_level(lz4f::CLEVEL_HIGH).build())?;
|
||||||
|
file.write_all(&chain[..]).await?;
|
||||||
|
file.shutdown().await?;
|
||||||
|
Ok(())
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Start the save loop for this chain
|
||||||
|
pub async fn host(state: State)
|
||||||
|
{
|
||||||
|
let to = &state.config().file;
|
||||||
|
let interval = state.config().save_interval();
|
||||||
|
while Arc::strong_count(state.when()) > 1 {
|
||||||
|
{
|
||||||
|
let chain = state.chain().read().await;
|
||||||
|
use std::ops::Deref;
|
||||||
|
if let Err(e) = save_now(chain.deref(), &to).await {
|
||||||
|
error!("Failed to save chain: {}", e);
|
||||||
|
} else {
|
||||||
|
info!("Saved chain to {:?}", to);
|
||||||
|
}
|
||||||
|
}
|
||||||
|
if state.has_shutdown() {
|
||||||
|
break;
|
||||||
|
}
|
||||||
|
OptionFuture::from(interval.map(|interval| time::delay_for(interval))).await;
|
||||||
|
state.when().notified().await;
|
||||||
|
}
|
||||||
|
trace!("Saver exiting");
|
||||||
|
}
|
||||||
|
|
||||||
|
/// Try to load a chain from this path
|
||||||
|
pub async fn load(from: impl AsRef<Path>) -> io::Result<Chain<String>>
|
||||||
|
{
|
||||||
|
debug!("Loading chain from {:?}", from.as_ref());
|
||||||
|
let file = OpenOptions::new()
|
||||||
|
.read(true)
|
||||||
|
.open(from).await?;
|
||||||
|
let mut whole = Vec::new();
|
||||||
|
let mut file = AsyncReadDecompressor::new(file)?;
|
||||||
|
tokio::io::copy(&mut file, &mut whole).await?;
|
||||||
|
serde_cbor::from_slice(&whole[..])
|
||||||
|
.map_err(|e| io::Error::new(io::ErrorKind::InvalidInput, e))
|
||||||
|
}
|
||||||
@ -0,0 +1,74 @@
|
|||||||
|
//! State
|
||||||
|
use super::*;
|
||||||
|
use tokio::{
|
||||||
|
sync::{
|
||||||
|
watch,
|
||||||
|
},
|
||||||
|
};
|
||||||
|
use config::Config;
|
||||||
|
|
||||||
|
#[derive(Debug, Clone)]
|
||||||
|
pub struct State
|
||||||
|
{
|
||||||
|
config: Arc<Config>, //to avoid cloning config
|
||||||
|
chain: Arc<RwLock<Chain<String>>>,
|
||||||
|
save: Arc<Notify>,
|
||||||
|
|
||||||
|
shutdown: Arc<watch::Sender<bool>>,
|
||||||
|
shutdown_recv: watch::Receiver<bool>,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl State
|
||||||
|
{
|
||||||
|
pub fn new(config: Config, chain: Arc<RwLock<Chain<String>>>, save: Arc<Notify>) -> Self
|
||||||
|
{
|
||||||
|
let (shutdown, shutdown_recv) = watch::channel(false);
|
||||||
|
Self {
|
||||||
|
config: Arc::new(config),
|
||||||
|
chain,
|
||||||
|
save,
|
||||||
|
shutdown: Arc::new(shutdown),
|
||||||
|
shutdown_recv,
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn config(&self) -> &Config
|
||||||
|
{
|
||||||
|
self.config.as_ref()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn notify_save(&self)
|
||||||
|
{
|
||||||
|
self.save.notify();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn chain(&self) -> &RwLock<Chain<String>>
|
||||||
|
{
|
||||||
|
&self.chain.as_ref()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn when(&self) -> &Arc<Notify>
|
||||||
|
{
|
||||||
|
&self.save
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn shutdown(self)
|
||||||
|
{
|
||||||
|
self.shutdown.broadcast(true).expect("Failed to communicate shutdown");
|
||||||
|
self.save.notify();
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn has_shutdown(&self) -> bool
|
||||||
|
{
|
||||||
|
*self.shutdown_recv.borrow()
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn on_shutdown(mut self)
|
||||||
|
{
|
||||||
|
if !self.has_shutdown() {
|
||||||
|
while let Some(false) = self.shutdown_recv.recv().await {
|
||||||
|
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
|
}
|
||||||
Loading…
Reference in new issue