rust/neon
1
0
Fork 0
mirror of https://github.com/neondatabase/neon.git synced 2026-05-29 19:10:38 +00:00
Code Issues Packages Projects Releases Wiki Activity

WIP

Browse Source
This commit is contained in:
Bojan Serafimov
2022-01-11 01:31:37 -05:00
parent 5b9391b51d
commit 9d6b78861d
3 changed files with 82 additions and 99 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
proxy/src/main.rs
View File

@@ -17,7 +17,8 @@ mod proxy;
mod state;
mod waiters;
fn main() -> anyhow::Result<()> {
#[tokio::main]
async fn main() -> anyhow::Result<()> {
let arg_matches = App::new("Zenith proxy/router")
.version(GIT_VERSION)
.arg(
@@ -97,20 +98,10 @@ fn main() -> anyhow::Result<()> {
println!("Starting mgmt on {}", state.conf.mgmt_address);
let mgmt_listener = tcp_listener::bind(state.conf.mgmt_address)?;
let threads = [
// Spawn a thread to listen for connections. It will spawn further threads
// for each connection.
thread::Builder::new()
.name("Listener thread".into())
.spawn(move || proxy::thread_main(state, pageserver_listener))?,
thread::Builder::new()
.name("Mgmt thread".into())
.spawn(move || mgmt::thread_main(state, mgmt_listener))?,
];
for t in threads {
t.join().unwrap()?;
}
tokio::try_join!(
proxy::thread_main(state, pageserver_listener),
mgmt::thread_main(state, mgmt_listener),
)?;
Ok(())
}

2
proxy/src/mgmt.rs
View File

@@ -16,7 +16,7 @@ use crate::{cplane_api::DatabaseInfo, ProxyState};
///
/// Listens for connections, and launches a new handler thread for each.
///
pub fn thread_main(state: &'static ProxyState, listener: TcpListener) -> anyhow::Result<()> {
pub async fn thread_main(state: &'static ProxyState, listener: TcpListener) -> anyhow::Result<()> {
loop {
let (socket, peer_addr) = listener.accept()?;
println!("accepted connection from {}", peer_addr);

158
proxy/src/proxy.rs
View File

@@ -14,6 +14,7 @@ use zenith_utils::postgres_backend::{self, PostgresBackend, ProtoState, Stream};
use zenith_utils::pq_proto::{BeMessage as Be, FeMessage as Fe, *};
use zenith_utils::sock_split::{ReadStream, WriteStream};
#[derive(Clone)]
struct CancelClosure {
socket_addr: SocketAddr,
cancel_token: tokio_postgres::CancelToken,
@@ -35,9 +36,14 @@ lazy_static! {
static ref CANCEL_MAP: Mutex<HashMap<CancelKeyData, CancelClosure>> = Mutex::new(HashMap::new());
}
thread_local! {
// Used to clean up the CANCEL_MAP. Might not be necessary if we use tokio thread pool in main loop.
static THREAD_CANCEL_KEY_DATA: Cell<Option<CancelKeyData>> = Cell::new(None);
/// Create new CancelKeyData with backend_pid that doesn't necessarily
/// correspond to the backend_pid of any actual backend.
fn fabricate_cancel_key_data() -> CancelKeyData {
let mut rng = StdRng::from_entropy();
CancelKeyData {
backend_pid: rng.gen(),
cancel_key: rng.gen(),
}
}
///
@@ -45,7 +51,7 @@ thread_local! {
///
/// Listens for connections, and launches a new handler thread for each.
///
pub fn thread_main(
pub async fn thread_main(
state: &'static ProxyState,
listener: std::net::TcpListener,
) -> anyhow::Result<()> {
@@ -54,23 +60,16 @@ pub fn thread_main(
println!("accepted connection from {}", peer_addr);
socket.set_nodelay(true).unwrap();
// TODO Use a threadpool instead. Maybe use tokio's threadpool by
// spawning a future into its runtime. Tokio's JoinError should
// allow us to handle cleanup properly even if the future panics.
thread::Builder::new()
.name("Proxy thread".into())
.spawn(move || {
if let Err(err) = proxy_conn_main(state, socket) {
println!("error: {}", err);
}
// Clean up CANCEL_MAP.
THREAD_CANCEL_KEY_DATA.with(|cell| {
if let Some(cancel_key_data) = cell.get() {
CANCEL_MAP.lock().remove(&cancel_key_data);
};
});
})?;
tokio::task::spawn(async move {
let cancel_key_data = fabricate_cancel_key_data();
let res = tokio::task::spawn(proxy_conn_main(state, socket, cancel_key_data)).await;
CANCEL_MAP.lock().remove(&cancel_key_data);
match res {
Err(join_err) => println!("join error: {}", join_err),
Ok(Err(conn_err)) => println!("connection error: {}", conn_err),
Ok(Ok(())) => {},
}
});
}
}
@@ -81,7 +80,7 @@ struct ProxyConnection {
pgb: PostgresBackend,
}
pub fn proxy_conn_main(state: &'static ProxyState, socket: TcpStream) -> anyhow::Result<()> {
pub async fn proxy_conn_main(state: &'static ProxyState, socket: TcpStream, cancel_key_data: CancelKeyData) -> anyhow::Result<()> {
let conn = ProxyConnection {
state,
psql_session_id: hex::encode(rand::random::<[u8; 8]>()),
@@ -93,7 +92,7 @@ pub fn proxy_conn_main(state: &'static ProxyState, socket: TcpStream) -> anyhow:
)?,
};
let (client, server) = match conn.handle_client()? {
let (client, server) = match conn.handle_client(cancel_key_data).await? {
Some(x) => x,
None => return Ok(()),
};
@@ -105,28 +104,41 @@ pub fn proxy_conn_main(state: &'static ProxyState, socket: TcpStream) -> anyhow:
_ => panic!("invalid stream type"),
};
proxy(client.split(), server.split())
proxy(client.split(), server.split()).await
}
impl ProxyConnection {
/// Returns Ok(None) when connection was successfully closed.
fn handle_client(mut self) -> anyhow::Result<Option<(Stream, TcpStream)>> {
let mut authenticate = || {
let (username, dbname) = match self.handle_startup()? {
Some(x) => x,
None => return Ok(None),
};
async fn handle_client(mut self, cancel_key_data: CancelKeyData) -> anyhow::Result<Option<(Stream, TcpStream)>> {
let (username, dbname) = match self.handle_startup().await? {
Some(x) => x,
None => return Ok(None),
};
// Both scenarios here should end up producing database credentials
if username.ends_with("@zenith") {
let dbinfo = {
if true || username.ends_with("@zenith") {
self.handle_existing_user(&username, &dbname).map(Some)
} else {
self.handle_new_user().map(Some)
}
};
let conn = match authenticate() {
Ok(Some(db_info)) => connect_to_db(db_info),
// let mut authenticate = || async {
// let (username, dbname) = match self.handle_startup().await? {
// Some(x) => x,
// None => return Ok(None),
// };
// // Both scenarios here should end up producing database credentials
// if true || username.ends_with("@zenith") {
// self.handle_existing_user(&username, &dbname).map(Some)
// } else {
// self.handle_new_user().map(Some)
// }
// };
let conn = match dbinfo {
Ok(Some(info)) => connect_to_db(info),
Ok(None) => return Ok(None),
Err(e) => {
// Report the error to the client
@@ -137,11 +149,8 @@ impl ProxyConnection {
// We'll get rid of this once migration to async is complete
let (pg_version, db_stream) = {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()?;
let (pg_version, stream, cancel_key_data) = runtime.block_on(conn)?;
let (pg_version, stream, cancel_closure) = conn.await?;
CANCEL_MAP.lock().insert(cancel_key_data, cancel_closure);
self.pgb
.write_message(&BeMessage::BackendKeyData(cancel_key_data))?;
let stream = stream.into_std()?;
@@ -161,7 +170,7 @@ impl ProxyConnection {
}
/// Returns Ok(None) when connection was successfully closed.
fn handle_startup(&mut self) -> anyhow::Result<Option<(String, String)>> {
async fn handle_startup(&mut self) -> anyhow::Result<Option<(String, String)>> {
let have_tls = self.pgb.tls_config.is_some();
let mut encrypted = false;
@@ -198,12 +207,9 @@ impl ProxyConnection {
return Ok(Some((get_param("user")?, get_param("database")?)));
}
FeStartupPacket::CancelRequest(cancel_key_data) => {
if let Some(cancel_closure) = CANCEL_MAP.lock().get(&cancel_key_data) {
let runtime = tokio::runtime::Builder::new_current_thread()
.enable_all()
.build()
.unwrap();
runtime.block_on(cancel_closure.try_cancel_query());
let entry = CANCEL_MAP.lock().get(&cancel_key_data).map(core::clone::Clone::clone);
if let Some(cancel_closure) = entry {
cancel_closure.try_cancel_query().await;
}
return Ok(None);
}
@@ -231,14 +237,21 @@ impl ProxyConnection {
.split_last()
.ok_or_else(|| anyhow!("unexpected password message"))?;
let cplane = CPlaneApi::new(&self.state.conf.auth_endpoint, &self.state.waiters);
let db_info = cplane.authenticate_proxy_request(
user,
db,
md5_response,
&md5_salt,
&self.psql_session_id,
)?;
let db_info = DatabaseInfo {
host: "localhost".into(),
port: 5432,
dbname: "postgres".into(),
user: "postgres".into(),
password: Some("postgres".into()),
};
// let cplane = CPlaneApi::new(&self.state.conf.auth_endpoint, &self.state.waiters);
// let db_info = cplane.authenticate_proxy_request(
// user,
// db,
// md5_response,
// &md5_salt,
// &self.psql_session_id,
// )?;
self.pgb
.write_message_noflush(&Be::AuthenticationOk)?
@@ -287,7 +300,7 @@ fn hello_message(redirect_uri: &str, session_id: &str) -> String {
/// Create a TCP connection to a postgres database, authenticate with it, and receive the ReadyForQuery message
async fn connect_to_db(
db_info: DatabaseInfo,
) -> anyhow::Result<(String, tokio::net::TcpStream, CancelKeyData)> {
) -> anyhow::Result<(String, tokio::net::TcpStream, CancelClosure)> {
// Make raw connection. When connect_raw finishes we've received ReadyForQuery.
let socket_addr = db_info.socket_addr()?;
let mut socket = tokio::net::TcpStream::connect(socket_addr).await?;
@@ -295,41 +308,21 @@ async fn connect_to_db(
// NOTE We effectively ignore some ParameterStatus and NoticeResponse
// messages here. Not sure if that could break something.
let (client, conn) = config.connect_raw(&mut socket, NoTls).await?;
// Save info for potentially cancelling the query later
let mut rng = StdRng::from_entropy();
let cancel_key_data = CancelKeyData {
// HACK We'd rather get the real backend_pid but tokio_postgres doesn't
// expose it and we don't want to do another roundtrip to query
// for it. The client will be able to notice that this is not the
// actual backend_pid, but backend_pid is not used for anything
// so it doesn't matter.
backend_pid: rng.gen(),
cancel_key: rng.gen(),
};
let cancel_closure = CancelClosure {
socket_addr,
cancel_token: client.cancel_token(),
};
CANCEL_MAP.lock().insert(cancel_key_data, cancel_closure);
THREAD_CANCEL_KEY_DATA.with(|cell| {
let prev_value = cell.replace(Some(cancel_key_data));
assert!(
prev_value.is_none(),
"THREAD_CANCEL_KEY_DATA was already set"
);
});
let version = conn.parameter("server_version").unwrap();
Ok((version.into(), socket, cancel_key_data))
Ok((version.into(), socket, cancel_closure))
}
/// Concurrently proxy both directions of the client and server connections
fn proxy(
async fn proxy(
(client_read, client_write): (ReadStream, WriteStream),
(server_read, server_write): (ReadStream, WriteStream),
) -> anyhow::Result<()> {
fn do_proxy(mut reader: impl io::Read, mut writer: WriteStream) -> io::Result<u64> {
async fn do_proxy(mut reader: impl io::Read, mut writer: WriteStream) -> io::Result<u64> {
/// FlushWriter will make sure that every message is sent as soon as possible
struct FlushWriter<W>(W);
@@ -354,10 +347,9 @@ fn proxy(
res
}
let client_to_server_jh = thread::spawn(move || do_proxy(client_read, server_write));
do_proxy(server_read, client_write)?;
client_to_server_jh.join().unwrap()?;
tokio::try_join!(
do_proxy(client_read, server_write),
do_proxy(server_read, client_write),
)?;
Ok(())
}