mnml/server/src/net.rs
2019-03-18 16:36:26 +11:00

97 lines
2.8 KiB
Rust

use tungstenite::Message;
use tungstenite::protocol::WebSocket;
use tungstenite::server::accept;
use tungstenite::error::Error;
use tungstenite::Message::Binary;
use std::net::{TcpListener, TcpStream};
use serde_cbor::{to_vec};
use std::env;
use std::thread::spawn;
use r2d2::{Pool};
use r2d2::{PooledConnection};
use r2d2_postgres::{TlsMode, PostgresConnectionManager};
static DB_POOL_SIZE: u32 = 20;
pub type Db = PooledConnection<PostgresConnectionManager>;
use rpc::{Rpc};
use util::{startup};
// struct Server {
// client: WebSocket<TcpStream>,
// rpc: Rpc,
// db: Pool<PostgresConnectionManager>,
// }
#[derive(Debug,Clone,Serialize,Deserialize)]
struct RpcErrorResponse {
err: String
}
fn receive(db: Db, rpc: &Rpc, msg: Message, client: &mut WebSocket<TcpStream>) -> Result<(), Error> {
match rpc.receive(msg, &db, client) {
Ok(reply) => {
let response = to_vec(&reply)
.expect("failed to serialize response");
client.write_message(Binary(response))
},
Err(e) => {
println!("{:?}", e);
let response = to_vec(&RpcErrorResponse { err: e.to_string() })
.expect("failed to serialize error response");
client.write_message(Binary(response))
}
}
}
pub fn db_connection(url: String) -> Pool<PostgresConnectionManager> {
let manager = PostgresConnectionManager::new(url, TlsMode::None)
.expect("could not instantiate pg manager");
Pool::builder()
.max_size(DB_POOL_SIZE)
.build(manager)
.expect("Failed to create pool.")
}
pub fn start() {
let database_url = env::var("DATABASE_URL")
.expect("DATABASE_URL must be set");
let pool = db_connection(database_url);
{
let startup_connection = pool.get().expect("unable to get db connection");
startup(startup_connection).unwrap();
}
let server = TcpListener::bind("0.0.0.0:40000").unwrap();
for stream in server.incoming() {
let db = pool.clone();
spawn (move || {
let mut websocket = accept(stream.unwrap()).unwrap();
let rpc = Rpc {};
loop {
match websocket.read_message() {
Ok(msg) => {
let db_connection = db.get().expect("unable to get db connection");
match receive(db_connection, &rpc, msg, &mut websocket) {
Ok(_r) => (),
Err(e) => println!("{:?}", e),
}
},
// connection is closed
Err(e) => {
println!("{:?}", e);
return;
}
};
}
});
}
}