magnumdb 0.4.8

High-performance embeddable database engine in Rust with B+ Tree storage, WAL durability, Volcano query engine, and Postgres wire protocol support.
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

use crate::network::pgwire::PgWireHandler;
use crate::sql::{Executor, Parser};
use crate::storage::Database;
use anyhow::Result;
use std::sync::Arc;
use std::time::Duration;
use tokio::io::{AsyncReadExt, AsyncWriteExt};
use tokio::net::TcpListener;
use tokio::sync::{RwLock, Semaphore};

pub struct Server {
    db: Arc<Database>,
    addr: String,
    max_connections: usize,
    idle_timeout: Duration,
}

impl Server {
    pub fn new(db: Database, addr: String) -> Self {
        Self {
            db: Arc::new(db),
            addr,
            max_connections: 128,
            idle_timeout: Duration::from_secs(30),
        }
    }

    /// Creates a server with custom connection limits.
    pub fn with_limits(mut self, max_connections: usize, idle_timeout_secs: u64) -> Self {
        self.max_connections = max_connections;
        self.idle_timeout = Duration::from_secs(idle_timeout_secs);
        self
    }

    pub async fn run(&self) -> Result<()> {
        let listener = TcpListener::bind(&self.addr).await?;
        println!("MagnumDB Server listening on {}", self.addr);

        let semaphore = Arc::new(Semaphore::new(self.max_connections));
        let idle_timeout = self.idle_timeout;

        loop {
            let (mut socket, peer) = listener.accept().await?;
            let db_ref = Arc::clone(&self.db);
            let permit = match semaphore.clone().try_acquire_owned() {
                Ok(permit) => permit,
                Err(_) => {
                    log::warn!("Connection limit reached, rejecting {}", peer);
                    let msg = b"Error: Too many connections\n";
                    let len = (msg.len() as u32).to_le_bytes();
                    let _ = socket.write_all(&len).await;
                    let _ = socket.write_all(msg).await;
                    continue;
                }
            };

            tokio::spawn(async move {
                let _permit = permit;
                log::info!("Client connected from {}", peer);
                let mut len_buf = [0u8; 4];

                loop {
                    let read_result = if idle_timeout.as_secs() > 0 {
                        tokio::time::timeout(idle_timeout, socket.read_exact(&mut len_buf)).await
                    } else {
                        Ok(socket.read_exact(&mut len_buf).await)
                    };

                    let read_result = match read_result {
                        Ok(r) => r,
                        Err(_) => {
                            log::info!("Client {} timed out after idle", peer);
                            break;
                        }
                    };

                    if read_result.is_err() {
                        break;
                    }

                    // Check for Postgres SSL Request / StartupMessage (big-endian length or SSL magic)
                    let be_len = i32::from_be_bytes(len_buf);
                    if (8..=10000).contains(&be_len) {
                        // Looks like PGWire client — delegate remaining handling to PgWireHandler
                        let pg_handler = PgWireHandler::new(Arc::clone(&db_ref));
                        // Re-assemble full socket stream with the initial len_buf already read
                        let _ = pg_handler.handle_connection(socket).await;
                        break;
                    }

                    let query_len = u32::from_le_bytes(len_buf) as usize;
                    if query_len > 16 * 1024 * 1024 {
                        log::warn!("Client {} sent oversized query ({}B), dropping", peer, query_len);
                        break;
                    }

                    let mut query_buf = vec![0u8; query_len];
                    if socket.read_exact(&mut query_buf).await.is_err() {
                        break;
                    }

                    let query = String::from_utf8_lossy(&query_buf);

                    let response = {
                        let mut executor = Executor::new(&db_ref);
                        match Parser::parse(&query) {
                            Ok(stmt) => match executor.execute(stmt) {
                                Ok(res) => res,
                                Err(e) => format!("Error: {}", e),
                            },
                            Err(e) => format!("Error: {}", e),
                        }
                    };

                    let resp_bytes = response.as_bytes();
                    let resp_len = (resp_bytes.len() as u32).to_le_bytes();

                    if socket.write_all(&resp_len).await.is_err() {
                        break;
                    }
                    if socket.write_all(resp_bytes).await.is_err() {
                        break;
                    }
                }
            });
        }
    }
}