Merge remote-tracking branch 'origin/main' into optimize-operation
This commit is contained in:
Yuriy Dupyn
commit
900608b2f4
17 changed files with 814 additions and 198 deletions
34
server/src/cancellation.rs
Normal file
34
server/src/cancellation.rs
Normal file
|
|
@ -0,0 +1,34 @@
|
|||
use std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicBool, Ordering};
|
||||
|
||||
pub struct ResetCancelToken {
|
||||
is_canceled: Arc<AtomicBool>,
|
||||
}
|
||||
|
||||
impl ResetCancelToken {
|
||||
pub fn new() -> Self {
|
||||
Self {
|
||||
is_canceled: Arc::new(AtomicBool::new(false)),
|
||||
}
|
||||
}
|
||||
|
||||
pub fn is_canceled(&self) -> bool {
|
||||
self.is_canceled.load(Ordering::SeqCst)
|
||||
}
|
||||
|
||||
pub fn cancel(&self) {
|
||||
self.is_canceled.store(true, Ordering::SeqCst);
|
||||
}
|
||||
|
||||
pub fn reset(&self) {
|
||||
self.is_canceled.store(false, Ordering::SeqCst);
|
||||
}
|
||||
}
|
||||
|
||||
impl Clone for ResetCancelToken {
|
||||
fn clone(&self) -> Self {
|
||||
Self {
|
||||
is_canceled: self.is_canceled.clone(),
|
||||
}
|
||||
}
|
||||
}
|
||||
28
server/src/config.rs
Normal file
28
server/src/config.rs
Normal file
|
|
@ -0,0 +1,28 @@
|
|||
use std::net::{IpAddr, Ipv4Addr, SocketAddr};
|
||||
use std::path::PathBuf;
|
||||
use clap::Parser;
|
||||
|
||||
const LOCAL_IPV4: IpAddr = IpAddr::V4(Ipv4Addr::LOCALHOST);
|
||||
|
||||
#[derive(Debug, Parser)]
|
||||
#[command(author, version, about)]
|
||||
pub struct Configuration {
|
||||
#[arg(short, long, default_value_t = LOCAL_IPV4, help = "IP address for the server to listen on")]
|
||||
address: IpAddr,
|
||||
#[arg(short, long, default_value = "5432", help = "Port for the server to listen on")]
|
||||
port: u16,
|
||||
#[arg(short, long, help = "Path to the data file")]
|
||||
file: PathBuf,
|
||||
}
|
||||
|
||||
impl Configuration {
|
||||
#[inline]
|
||||
pub fn get_socket_address(&self) -> SocketAddr {
|
||||
SocketAddr::new(self.address, self.port)
|
||||
}
|
||||
|
||||
#[inline]
|
||||
pub fn get_file_path(&self) -> &PathBuf {
|
||||
&self.file
|
||||
}
|
||||
}
|
||||
|
|
@ -1,78 +1,139 @@
|
|||
use minisql::interpreter::State;
|
||||
use parser::{parse_and_validate, Error};
|
||||
use std::collections::HashMap;
|
||||
use std::sync::Arc;
|
||||
|
||||
use clap::Parser;
|
||||
use tokio::io::{BufReader, BufWriter};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
use tokio::sync::{Mutex, RwLock};
|
||||
|
||||
use minisql::interpreter::{Response, State};
|
||||
use parser::parse_and_validate;
|
||||
use proto::handshake::errors::ServerHandshakeError;
|
||||
use proto::handshake::request::HandshakeRequest;
|
||||
use proto::handshake::response::HandshakeResponse;
|
||||
use proto::handshake::server::do_server_handshake;
|
||||
use proto::message::backend::{
|
||||
BackendMessage, ColumnDescription, CommandCompleteData, DataRowData, ErrorResponseData,
|
||||
ReadyForQueryData, RowDescriptionData,
|
||||
};
|
||||
use proto::message::frontend::FrontendMessage;
|
||||
use proto::reader::oneway::OneWayProtoReader;
|
||||
use proto::reader::frontend::FrontendProtoReader;
|
||||
use proto::reader::protoreader::ProtoReader;
|
||||
use proto::writer::backend::BackendProtoWriter;
|
||||
use proto::writer::protowriter::{ProtoFlush, ProtoWriter};
|
||||
use tokio::io::{BufReader, BufWriter};
|
||||
use tokio::net::{TcpListener, TcpStream};
|
||||
|
||||
use crate::cancellation::ResetCancelToken;
|
||||
use crate::config::Configuration;
|
||||
use crate::proto_wrapper::{CompleteStatus, ServerProto};
|
||||
|
||||
mod config;
|
||||
mod proto_wrapper;
|
||||
mod cancellation;
|
||||
|
||||
type TokenStore = Arc<Mutex<HashMap<(i32, i32), ResetCancelToken>>>;
|
||||
type SharedDbState = Arc<RwLock<State>>;
|
||||
|
||||
#[tokio::main]
|
||||
async fn main() -> anyhow::Result<()> {
|
||||
let addr = "0.0.0.0:5432";
|
||||
let config = Configuration::parse();
|
||||
|
||||
let state = Arc::new(RwLock::new(State::new()));
|
||||
let tokens = Arc::new(Mutex::new(HashMap::<(i32, i32), ResetCancelToken>::new()));
|
||||
|
||||
let addr = config.get_socket_address();
|
||||
let listener = TcpListener::bind(&addr).await?;
|
||||
println!("Server started at {addr}");
|
||||
|
||||
loop {
|
||||
let state = state.clone();
|
||||
let tokens = tokens.clone();
|
||||
|
||||
let (socket, _) = listener.accept().await?;
|
||||
println!("New client connected: {}", socket.peer_addr()?);
|
||||
tokio::spawn(async move {
|
||||
let reason = handle_stream(socket).await;
|
||||
let reason = handle_stream(socket, state, tokens).await;
|
||||
println!("Client disconnected: {reason:?}");
|
||||
});
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_stream(mut stream: TcpStream) -> anyhow::Result<()> {
|
||||
async fn handle_stream(mut stream: TcpStream, state: SharedDbState, tokens: TokenStore) -> anyhow::Result<()> {
|
||||
let (reader, writer) = stream.split();
|
||||
let mut writer = ProtoWriter::new(BufWriter::new(writer));
|
||||
let mut reader = ProtoReader::new(BufReader::new(reader), 1024);
|
||||
|
||||
let response = HandshakeResponse::new("minisql", 123, 123);
|
||||
// Create a token with random PID and key
|
||||
let (pid, key, token) = create_token(&tokens).await?;
|
||||
|
||||
let request = do_server_handshake(&mut writer, &mut reader, response).await?;
|
||||
// Handle handshake
|
||||
let response = HandshakeResponse::new("minisql", pid, key);
|
||||
let request = do_server_handshake(&mut writer, &mut reader, response).await;
|
||||
|
||||
println!("Handshake complete:\n{request:?}");
|
||||
let mut state = State::new();
|
||||
let result = match request {
|
||||
Ok(req) => handle_connection(&mut reader, &mut writer, req, state, token).await,
|
||||
Err(ServerHandshakeError::IsCancelRequest(cancel)) => handle_cancellation(cancel.pid, cancel.secret, &tokens).await,
|
||||
Err(e) => Err(anyhow::anyhow!("Error during handshake: {:?}", e)),
|
||||
};
|
||||
|
||||
// Release cancellation token
|
||||
let mut tokens = tokens.lock().await;
|
||||
tokens.remove(&(pid, key));
|
||||
|
||||
result
|
||||
}
|
||||
|
||||
fn random_pid_key() -> (i32, i32) {
|
||||
let pid = rand::random::<i32>();
|
||||
let key = rand::random::<i32>();
|
||||
(pid, key)
|
||||
}
|
||||
|
||||
async fn create_token(tokens: &TokenStore) -> anyhow::Result<(i32, i32, ResetCancelToken)> {
|
||||
let token = ResetCancelToken::new();
|
||||
let mut tokens = tokens.lock().await;
|
||||
loop {
|
||||
let pid_key = random_pid_key();
|
||||
if !tokens.contains_key(&pid_key) {
|
||||
tokens.insert(pid_key, token.clone());
|
||||
|
||||
let (pid, key) = pid_key;
|
||||
return Ok((pid, key, token));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
async fn handle_cancellation(pid: i32, key: i32, tokens: &TokenStore) -> anyhow::Result<()> {
|
||||
println!("Cancel request, PID: {}, Key: {}", pid, key);
|
||||
|
||||
let tokens = tokens.lock().await;
|
||||
let token = tokens.get(&(pid, key));
|
||||
match token {
|
||||
Some(t) => t.cancel(),
|
||||
None => return Err(anyhow::anyhow!("Invalid PID and Key cancel combination")),
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn handle_connection<R, W>(reader: &mut R, writer: &mut W, request: HandshakeRequest, state: SharedDbState, token: ResetCancelToken) -> anyhow::Result<()>
|
||||
where
|
||||
R: FrontendProtoReader + Send,
|
||||
W: BackendProtoWriter + ProtoFlush + Send,
|
||||
{
|
||||
println!("Client connected: {:?}", request);
|
||||
|
||||
loop {
|
||||
println!("Waiting for next message");
|
||||
let next: FrontendMessage = reader.read_proto().await?;
|
||||
|
||||
match next {
|
||||
FrontendMessage::Terminate => {
|
||||
println!("Received Terminate");
|
||||
break;
|
||||
}
|
||||
FrontendMessage::Query(data) => {
|
||||
println!("Received Query: {:?}", data);
|
||||
let db_schema = state.db_schema();
|
||||
match parse_and_validate(data.query.as_str().to_string(), &db_schema) {
|
||||
Ok(operation) => {
|
||||
match state.interpret(operation) {
|
||||
Ok(_) => {
|
||||
send_query_response(&mut writer).await?;
|
||||
}
|
||||
Err(err) => {
|
||||
send_error_response(&mut writer, &format!("error interpreting: {:?}", err)).await?;
|
||||
}
|
||||
}
|
||||
},
|
||||
Err(Error::ParsingError(err)) => {
|
||||
send_error_response(&mut writer, &format!("parsing error: {:?}", err)).await?;
|
||||
let result = handle_query(writer, &state, data.query.into(), &token).await;
|
||||
match result {
|
||||
Ok(_) => {}
|
||||
Err(e) => {
|
||||
writer.write_error_message(&e.to_string()).await?
|
||||
}
|
||||
Err(Error::ValidationError(v)) => {
|
||||
send_error_response(&mut writer, &format!("validation error: {:?}", v)).await?;
|
||||
}
|
||||
};
|
||||
send_ready_for_query(&mut writer).await?;
|
||||
}
|
||||
writer.write_ready_for_query().await?;
|
||||
}
|
||||
}
|
||||
writer.flush().await?;
|
||||
|
|
@ -81,117 +142,47 @@ async fn handle_stream(mut stream: TcpStream) -> anyhow::Result<()> {
|
|||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_error_response(
|
||||
writer: &mut impl BackendProtoWriter,
|
||||
error_message: &str,
|
||||
) -> anyhow::Result<()> {
|
||||
writer
|
||||
.write_proto(
|
||||
ErrorResponseData {
|
||||
code: b'M',
|
||||
message: error_message.to_string().into(),
|
||||
async fn handle_query<W>(writer: &mut W, state: &SharedDbState, query: String, token: &ResetCancelToken) -> anyhow::Result<()>
|
||||
where
|
||||
W: BackendProtoWriter + ProtoFlush + Send,
|
||||
{
|
||||
let operation = {
|
||||
let state = state.read().await;
|
||||
let db_schema = state.db_schema();
|
||||
parse_and_validate(query, &db_schema)?
|
||||
};
|
||||
|
||||
let mut state = state.write().await;
|
||||
let response = state.interpret(operation)?;
|
||||
|
||||
match response {
|
||||
Response::Deleted(i) => writer.write_command_complete(CompleteStatus::Delete(i)).await?,
|
||||
Response::Inserted => writer.write_command_complete(CompleteStatus::Insert { oid: 0, rows: 1 }).await?,
|
||||
Response::Selected(schema, mut rows) => {
|
||||
match rows.next() {
|
||||
Some(row) => {
|
||||
writer.write_table_header(&schema, &row).await?;
|
||||
writer.write_table_row(&row).await?;
|
||||
|
||||
let mut sent_rows = 1;
|
||||
for row in rows {
|
||||
sent_rows += 1;
|
||||
writer.write_table_row(&row).await?;
|
||||
if token.is_canceled() {
|
||||
token.reset();
|
||||
break;
|
||||
}
|
||||
}
|
||||
|
||||
writer.write_command_complete(CompleteStatus::Select(sent_rows)).await?;
|
||||
}
|
||||
_ => {
|
||||
writer.write_command_complete(CompleteStatus::Select(0)).await?;
|
||||
}
|
||||
}
|
||||
.into(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_ready_for_query(writer: &mut impl BackendProtoWriter) -> anyhow::Result<()> {
|
||||
writer
|
||||
.write_proto(BackendMessage::from(ReadyForQueryData { status: b'I' }))
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_empty_query(writer: &mut impl BackendProtoWriter) -> anyhow::Result<()> {
|
||||
writer
|
||||
.write_proto(BackendMessage::EmptyQueryResponse)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_row_description(writer: &mut impl BackendProtoWriter) -> anyhow::Result<()> {
|
||||
let columns = vec![
|
||||
ColumnDescription {
|
||||
name: "id".to_string().into(),
|
||||
table_oid: 123,
|
||||
column_index: 1,
|
||||
type_oid: 23,
|
||||
type_size: 4,
|
||||
type_modifier: -1,
|
||||
format_code: 0,
|
||||
},
|
||||
ColumnDescription {
|
||||
name: "argument".to_string().into(),
|
||||
table_oid: 123,
|
||||
column_index: 2,
|
||||
type_oid: 23,
|
||||
type_size: 4,
|
||||
type_modifier: -1,
|
||||
format_code: 0,
|
||||
},
|
||||
ColumnDescription {
|
||||
name: "description".to_string().into(),
|
||||
table_oid: 123,
|
||||
column_index: 3,
|
||||
type_oid: 1043,
|
||||
type_size: 32,
|
||||
type_modifier: -1,
|
||||
format_code: 0,
|
||||
},
|
||||
];
|
||||
|
||||
writer
|
||||
.write_proto(
|
||||
RowDescriptionData {
|
||||
columns: columns.into(),
|
||||
}
|
||||
.into(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn send_query_response(writer: &mut impl BackendProtoWriter) -> anyhow::Result<()> {
|
||||
send_row_description(writer).await?;
|
||||
|
||||
write_row(writer, b"0", b"1337", b"auto").await?;
|
||||
write_row(writer, b"1", b"69", b"bus").await?;
|
||||
write_row(writer, b"2", b"420", b"kolo").await?;
|
||||
|
||||
writer
|
||||
.write_proto(
|
||||
CommandCompleteData {
|
||||
tag: "SELECT 3".to_string().into(),
|
||||
}
|
||||
.into(),
|
||||
)
|
||||
.await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_row(
|
||||
writer: &mut impl BackendProtoWriter,
|
||||
first: &[u8],
|
||||
second: &[u8],
|
||||
third: &[u8],
|
||||
) -> anyhow::Result<()> {
|
||||
let row_data = vec![
|
||||
first.to_vec().into(),
|
||||
second.to_vec().into(),
|
||||
third.to_vec().into(),
|
||||
]
|
||||
.into();
|
||||
|
||||
writer
|
||||
.write_proto(DataRowData { columns: row_data }.into())
|
||||
.await?;
|
||||
}
|
||||
_ => {}
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
|
|
|||
104
server/src/proto_wrapper.rs
Normal file
104
server/src/proto_wrapper.rs
Normal file
|
|
@ -0,0 +1,104 @@
|
|||
use async_trait::async_trait;
|
||||
use minisql::restricted_row::RestrictedRow;
|
||||
use minisql::schema::TableSchema;
|
||||
use minisql::type_system::{Value};
|
||||
use proto::message::backend::{BackendMessage, ColumnDescription, CommandCompleteData, DataRowData, ErrorResponseData, ReadyForQueryData, RowDescriptionData};
|
||||
use proto::message::primitive::pglist::PgList;
|
||||
use proto::writer::backend::BackendProtoWriter;
|
||||
|
||||
pub enum CompleteStatus {
|
||||
Insert {
|
||||
oid: i32,
|
||||
rows: i32,
|
||||
},
|
||||
Delete(usize),
|
||||
Select(usize),
|
||||
}
|
||||
|
||||
impl CompleteStatus {
|
||||
fn to_string(&self) -> String {
|
||||
match self {
|
||||
CompleteStatus::Insert { oid, rows } => format!("INSERT {} {}", oid, rows),
|
||||
CompleteStatus::Delete(rows) => format!("DELETE {}", rows),
|
||||
CompleteStatus::Select(rows) => format!("SELECT {}", rows),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
pub trait ServerProto {
|
||||
async fn write_error_message(&mut self, error_message: &str) -> anyhow::Result<()>;
|
||||
async fn write_ready_for_query(&mut self) -> anyhow::Result<()>;
|
||||
async fn write_empty_query(&mut self) -> anyhow::Result<()>;
|
||||
async fn write_table_header(&mut self, table_schema: &TableSchema, row: &RestrictedRow) -> anyhow::Result<()>;
|
||||
async fn write_table_row(&mut self, row: &RestrictedRow) -> anyhow::Result<()>;
|
||||
async fn write_command_complete(&mut self, status: CompleteStatus) -> anyhow::Result<()>;
|
||||
}
|
||||
|
||||
#[async_trait]
|
||||
impl<W> ServerProto for W where W: BackendProtoWriter + Send {
|
||||
async fn write_error_message(&mut self, error_message: &str) -> anyhow::Result<()> {
|
||||
self.write_proto(ErrorResponseData {
|
||||
code: b'M',
|
||||
message: format!("{error_message}\0").into(),
|
||||
}.into()).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_ready_for_query(&mut self) -> anyhow::Result<()> {
|
||||
self.write_proto(ReadyForQueryData { status: b'I' }.into()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_empty_query(&mut self) -> anyhow::Result<()> {
|
||||
self.write_proto(BackendMessage::EmptyQueryResponse).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_table_header(&mut self, table_schema: &TableSchema, row: &RestrictedRow) -> anyhow::Result<()> {
|
||||
let columns = row.iter()
|
||||
.map(|(index, value)| value_to_column_description(table_schema, value, index))
|
||||
.collect::<anyhow::Result<Vec<ColumnDescription>>>()?;
|
||||
|
||||
self.write_proto(RowDescriptionData { columns: columns.into() }.into()).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_table_row(&mut self, row: &RestrictedRow) -> anyhow::Result<()> {
|
||||
let values = row.iter()
|
||||
.map(|(_, value)| value.as_text_bytes().into())
|
||||
.collect::<Vec<PgList<u8, i32>>>();
|
||||
|
||||
self.write_proto(BackendMessage::DataRow(DataRowData {
|
||||
columns: values.into(),
|
||||
})).await?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
async fn write_command_complete(&mut self, status: CompleteStatus) -> anyhow::Result<()> {
|
||||
self.write_proto(BackendMessage::CommandComplete(CommandCompleteData {
|
||||
tag: status.to_string().into(),
|
||||
})).await?;
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
fn value_to_column_description(schema: &TableSchema, value: &Value, index: &usize) -> anyhow::Result<ColumnDescription> {
|
||||
let name = schema.column_name_from_column_position(*index)?;
|
||||
|
||||
let table_oid = schema.table_name().as_bytes().as_ptr() as i32;
|
||||
let column_index = (*index).try_into()?;
|
||||
let type_oid = value.type_oid();
|
||||
let type_size = value.type_size();
|
||||
|
||||
Ok(ColumnDescription {
|
||||
name: name.to_string().into(),
|
||||
table_oid,
|
||||
column_index,
|
||||
type_oid,
|
||||
type_size,
|
||||
type_modifier: -1,
|
||||
format_code: 0, // text format
|
||||
})
|
||||
}
|
||||
Loading…
Add table
Add a link
Reference in a new issue