omedusyo/minisql
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Merge remote-tracking branch 'origin/main' into optimize-operation

Browse source
This commit is contained in:
Yuriy Dupyn 2024-01-28 00:26:56 +01:00
commit 900608b2f4
17 changed files with 814 additions and 198 deletions
Download patch file Download diff file Expand all files Collapse all files

7
server/Cargo.toml
View file

@ -8,6 +8,9 @@ edition = "2021"
[dependencies]
tokio = { version = "1.35.1", features = ["full"] }
anyhow = "1.0.76"
proto = { path = "../proto" }
clap = { version = "4.4.18", features = ["derive"] }
async-trait = "0.1.74"
rand = "0.8.5"
minisql = { path = "../minisql" }
parser = { path = "../parser" }
proto = { path = "../proto" }
parser = { path = "../parser" }

34
server/src/cancellation.rs Normal file
View 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
View 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
}
}

287
server/src/main.rs
View 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
View 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
})
}