rrdb 0.0.2

ready
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

use std::sync::Arc;

use crate::errors::execute_error::ExecuteError;
use crate::errors::RRDBError;
use crate::executor::config::global::GlobalConfig;
use crate::executor::predule::Executor;
use crate::logger::predule::Logger;
use crate::pgwire::predule::Connection;
use crate::server::channel::ChannelResponse;
use crate::server::predule::{ChannelRequest, SharedState};

use futures::future::join_all;
use tokio::net::TcpListener;
use tokio::sync::mpsc;

use super::client::ClientInfo;

pub struct Server {
    pub config: Arc<GlobalConfig>,
}

impl Server {
    pub fn new(config: GlobalConfig) -> Self {
        Self {
            config: Arc::new(config),
        }
    }

    /// 메인 서버 루프.
    /// 여러개의 태스크 제어
    pub async fn run(&self) -> Result<(), RRDBError> {
        // TODO: 인덱스 로딩 등 기본 로직 실행.

        let (request_sender, mut request_receiver) = mpsc::channel::<ChannelRequest>(1000);

        // background task
        // 쿼리 실행 요청을 전달받음
        let config = self.config.clone();

        let background_task = tokio::spawn(async move {
            while let Some(request) = request_receiver.recv().await {
                let config = config.clone();

                // 쿼리 실행 태스크
                tokio::spawn(async move {
                    let executor = Executor::new(config);
                    let result = executor
                        .process_query(request.statement, request.connection_id)
                        .await;

                    match result {
                        Ok(result) => {
                            if let Err(_response) = request
                                .response_sender
                                .send(ChannelResponse { result: Ok(result) })
                            {
                                Logger::error("channel send failed");
                            }
                        }
                        Err(error) => {
                            let error = error.to_string();
                            if let Err(_response) = request.response_sender.send(ChannelResponse {
                                result: Err(ExecuteError::wrap(error)),
                            }) {
                                Logger::error("channel send failed");
                            }
                        }
                    }
                });
            }
        });

        // connection task
        // client와의 커넥션 처리 루프
        let listener = TcpListener::bind((self.config.host.to_owned(), self.config.port as u16))
            .await
            .map_err(|error| ExecuteError::wrap(error.to_string()))?;

        let config = self.config.clone();
        let connection_task = tokio::spawn(async move {
            loop {
                let accepted = listener.accept().await;

                let (stream, address) = match accepted {
                    Ok((stream, address)) => (stream, address),
                    Err(error) => {
                        Logger::error(format!("socket error {:?}", error));
                        continue;
                    }
                };

                let client_info = ClientInfo {
                    ip: address.ip(),
                    connection_id: uuid::Uuid::new_v4().to_string(),
                    database: "None".into(),
                };

                let shared_state = SharedState {
                    sender: request_sender.clone(),
                    client_info,
                };

                let config = config.clone();
                tokio::spawn(async move {
                    let mut conn = Connection::new(shared_state, config);
                    if let Err(error) = conn.run(stream).await {
                        Logger::error(format!("connection error {:?}", error));
                    }
                });
            }
        });

        Logger::info(format!(
            "Server is running on {}:{}",
            self.config.host, self.config.port
        ));

        join_all(vec![connection_task, background_task]).await;

        Ok(())
    }
}