parent
35f1de6c5e
commit
ecc8854e44
GPG Key ID:
284488987C31F630
@ -1,2 +1,3 @@
|
||||
/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