From 1fa4f93765415c8003f19e47fb771490087fe20f Mon Sep 17 00:00:00 2001 From: Bugen Zhao Date: Wed, 22 Jan 2025 17:07:51 +0800 Subject: [PATCH] refactor(frontend): use empty tcp stream for frontend health check in docker-compose (#20219) Signed-off-by: Bugen Zhao --- Cargo.lock | 13 +++++++ docker/docker-compose-distributed.yml | 2 +- docker/docker-compose-with-azblob.yml | 2 +- docker/docker-compose-with-gcs.yml | 2 +- docker/docker-compose-with-hdfs.yml | 2 +- docker/docker-compose-with-local-fs.yml | 2 +- docker/docker-compose-with-obs.yml | 2 +- docker/docker-compose-with-oss.yml | 2 +- docker/docker-compose-with-s3.yml | 2 +- docker/docker-compose-with-sqlite.yml | 2 +- docker/docker-compose.yml | 2 +- .../docker/hive/docker-compose.yml | 4 +-- .../docker/jdbc/docker-compose.yml | 4 +-- .../docker/rest/docker-compose.yml | 4 +-- .../docker/storage/docker-compose.yml | 4 +-- .../docker/hive/docker-compose.yml | 4 +-- .../docker/jdbc/docker-compose.yml | 4 +-- .../docker/rest/docker-compose.yml | 4 +-- .../docker/storage/docker-compose.yml | 4 +-- src/utils/pgwire/Cargo.toml | 1 + src/utils/pgwire/src/pg_message.rs | 36 +++++++------------ 21 files changed, 52 insertions(+), 50 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index b1f28fdb8185c..1f45fcad8c7a5 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -8781,6 +8781,18 @@ dependencies = [ "hmac", ] +[[package]] +name = "peekable" +version = "0.2.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8a57300da5f2129ddfba13137c21ff56b1755da8c68c77576c052ec96934c49b" +dependencies = [ + "bytes", + "pin-project-lite", + "smallvec", + "tokio", +] + [[package]] name = "pem" version = "3.0.2" @@ -8855,6 +8867,7 @@ dependencies = [ "openssl", "panic-message", "parking_lot 0.12.1", + "peekable", "postgres-types", "reqwest 0.12.4", "risingwave_common", diff --git a/docker/docker-compose-distributed.yml b/docker/docker-compose-distributed.yml index b1aac3e5b4a87..4e4a04f95d575 100644 --- a/docker/docker-compose-distributed.yml +++ b/docker/docker-compose-distributed.yml @@ -135,7 +135,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 5s retries: 5 diff --git a/docker/docker-compose-with-azblob.yml b/docker/docker-compose-with-azblob.yml index 8166012c244dc..ecafec709da1a 100644 --- a/docker/docker-compose-with-azblob.yml +++ b/docker/docker-compose-with-azblob.yml @@ -68,7 +68,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose-with-gcs.yml b/docker/docker-compose-with-gcs.yml index 9640ef6fc624b..25bcfe1657062 100644 --- a/docker/docker-compose-with-gcs.yml +++ b/docker/docker-compose-with-gcs.yml @@ -68,7 +68,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose-with-hdfs.yml b/docker/docker-compose-with-hdfs.yml index 7b69bdf3d29bc..54dc8124d3be5 100644 --- a/docker/docker-compose-with-hdfs.yml +++ b/docker/docker-compose-with-hdfs.yml @@ -108,7 +108,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 5s retries: 5 diff --git a/docker/docker-compose-with-local-fs.yml b/docker/docker-compose-with-local-fs.yml index 3cb3c3ec3cbc0..b9220b2c161a6 100644 --- a/docker/docker-compose-with-local-fs.yml +++ b/docker/docker-compose-with-local-fs.yml @@ -67,7 +67,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose-with-obs.yml b/docker/docker-compose-with-obs.yml index 0fa7e5529ea0b..facf56d4fba7c 100644 --- a/docker/docker-compose-with-obs.yml +++ b/docker/docker-compose-with-obs.yml @@ -68,7 +68,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose-with-oss.yml b/docker/docker-compose-with-oss.yml index 436e4fbd8d892..30b6f2abcbae9 100644 --- a/docker/docker-compose-with-oss.yml +++ b/docker/docker-compose-with-oss.yml @@ -68,7 +68,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose-with-s3.yml b/docker/docker-compose-with-s3.yml index 834639f2b90e3..8582922ddddd2 100644 --- a/docker/docker-compose-with-s3.yml +++ b/docker/docker-compose-with-s3.yml @@ -68,7 +68,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose-with-sqlite.yml b/docker/docker-compose-with-sqlite.yml index bc96966b5cf2a..f73eae95214e7 100644 --- a/docker/docker-compose-with-sqlite.yml +++ b/docker/docker-compose-with-sqlite.yml @@ -68,7 +68,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/docker/docker-compose.yml b/docker/docker-compose.yml index 5509de864ed6f..7e9a6aea5fb8d 100644 --- a/docker/docker-compose.yml +++ b/docker/docker-compose.yml @@ -70,7 +70,7 @@ services: - CMD-SHELL - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/6660; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5688; exit $$?;' - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/5690; exit $$?;' interval: 1s timeout: 5s diff --git a/integration_tests/iceberg-sink2/docker/hive/docker-compose.yml b/integration_tests/iceberg-sink2/docker/hive/docker-compose.yml index 449d6043e6e95..71a3b03ac13c3 100644 --- a/integration_tests/iceberg-sink2/docker/hive/docker-compose.yml +++ b/integration_tests/iceberg-sink2/docker/hive/docker-compose.yml @@ -47,7 +47,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s environment: @@ -113,4 +113,4 @@ volumes: networks: iceberg_net: - name: iceberg \ No newline at end of file + name: iceberg diff --git a/integration_tests/iceberg-sink2/docker/jdbc/docker-compose.yml b/integration_tests/iceberg-sink2/docker/jdbc/docker-compose.yml index 714fcdf6e0b95..0e63ef6e74e69 100644 --- a/integration_tests/iceberg-sink2/docker/jdbc/docker-compose.yml +++ b/integration_tests/iceberg-sink2/docker/jdbc/docker-compose.yml @@ -48,7 +48,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s environment: @@ -93,4 +93,4 @@ volumes: networks: iceberg_net: - name: iceberg \ No newline at end of file + name: iceberg diff --git a/integration_tests/iceberg-sink2/docker/rest/docker-compose.yml b/integration_tests/iceberg-sink2/docker/rest/docker-compose.yml index ee3e1da9b62da..120d84b9e361e 100644 --- a/integration_tests/iceberg-sink2/docker/rest/docker-compose.yml +++ b/integration_tests/iceberg-sink2/docker/rest/docker-compose.yml @@ -55,7 +55,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s networks: @@ -96,4 +96,4 @@ volumes: external: false networks: - iceberg_net: \ No newline at end of file + iceberg_net: diff --git a/integration_tests/iceberg-sink2/docker/storage/docker-compose.yml b/integration_tests/iceberg-sink2/docker/storage/docker-compose.yml index 4e0ec11eff500..bf8a1d8f7d669 100644 --- a/integration_tests/iceberg-sink2/docker/storage/docker-compose.yml +++ b/integration_tests/iceberg-sink2/docker/storage/docker-compose.yml @@ -31,7 +31,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s networks: @@ -72,4 +72,4 @@ volumes: external: false networks: - iceberg_net: \ No newline at end of file + iceberg_net: diff --git a/integration_tests/iceberg-source/docker/hive/docker-compose.yml b/integration_tests/iceberg-source/docker/hive/docker-compose.yml index 449d6043e6e95..71a3b03ac13c3 100644 --- a/integration_tests/iceberg-source/docker/hive/docker-compose.yml +++ b/integration_tests/iceberg-source/docker/hive/docker-compose.yml @@ -47,7 +47,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s environment: @@ -113,4 +113,4 @@ volumes: networks: iceberg_net: - name: iceberg \ No newline at end of file + name: iceberg diff --git a/integration_tests/iceberg-source/docker/jdbc/docker-compose.yml b/integration_tests/iceberg-source/docker/jdbc/docker-compose.yml index 714fcdf6e0b95..0e63ef6e74e69 100644 --- a/integration_tests/iceberg-source/docker/jdbc/docker-compose.yml +++ b/integration_tests/iceberg-source/docker/jdbc/docker-compose.yml @@ -48,7 +48,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s environment: @@ -93,4 +93,4 @@ volumes: networks: iceberg_net: - name: iceberg \ No newline at end of file + name: iceberg diff --git a/integration_tests/iceberg-source/docker/rest/docker-compose.yml b/integration_tests/iceberg-source/docker/rest/docker-compose.yml index ee3e1da9b62da..120d84b9e361e 100644 --- a/integration_tests/iceberg-source/docker/rest/docker-compose.yml +++ b/integration_tests/iceberg-source/docker/rest/docker-compose.yml @@ -55,7 +55,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s networks: @@ -96,4 +96,4 @@ volumes: external: false networks: - iceberg_net: \ No newline at end of file + iceberg_net: diff --git a/integration_tests/iceberg-source/docker/storage/docker-compose.yml b/integration_tests/iceberg-source/docker/storage/docker-compose.yml index 4e0ec11eff500..bf8a1d8f7d669 100644 --- a/integration_tests/iceberg-source/docker/storage/docker-compose.yml +++ b/integration_tests/iceberg-source/docker/storage/docker-compose.yml @@ -31,7 +31,7 @@ services: healthcheck: test: - CMD-SHELL - - bash -c 'printf \"GET / HTTP/1.1\n\n\" > /dev/tcp/127.0.0.1/4566; exit $$?;' + - bash -c '> /dev/tcp/127.0.0.1/4566; exit $$?;' interval: 1s timeout: 30s networks: @@ -72,4 +72,4 @@ volumes: external: false networks: - iceberg_net: \ No newline at end of file + iceberg_net: diff --git a/src/utils/pgwire/Cargo.toml b/src/utils/pgwire/Cargo.toml index 2275ed5a20571..76074e78fa442 100644 --- a/src/utils/pgwire/Cargo.toml +++ b/src/utils/pgwire/Cargo.toml @@ -19,6 +19,7 @@ jsonwebtoken = "9" openssl = "0.10.66" panic-message = "0.3" parking_lot = { workspace = true } +peekable = { version = "0.2", features = ["tokio"] } postgres-types = { version = "0.2.6" } reqwest = "0.12.2" risingwave_common = { workspace = true } diff --git a/src/utils/pgwire/src/pg_message.rs b/src/utils/pgwire/src/pg_message.rs index 363ca4e53e92e..a932356c5c2e4 100644 --- a/src/utils/pgwire/src/pg_message.rs +++ b/src/utils/pgwire/src/pg_message.rs @@ -17,9 +17,10 @@ use std::ffi::CStr; use std::io::{Error, ErrorKind, IoSlice, Result, Write}; use anyhow::anyhow; -use byteorder::{BigEndian, ByteOrder, NetworkEndian}; +use byteorder::{BigEndian, ByteOrder}; /// Part of code learned from . use bytes::{Buf, BufMut, Bytes, BytesMut}; +use peekable::tokio::AsyncPeekable; use tokio::io::{AsyncRead, AsyncReadExt}; use crate::error_or_notice::ErrorOrNoticeMessage; @@ -297,31 +298,18 @@ impl FeMessage { impl FeStartupMessage { /// Read startup message from the stream. pub async fn read(stream: &mut (impl AsyncRead + Unpin)) -> Result { - let mut buffer1 = vec![0; 1]; - let result = stream.read_exact(&mut buffer1).await; - let filled1 = match result { - Ok(n) => n, - Err(err) => { - // Detect whether it is a health check. - if err.kind() == ErrorKind::UnexpectedEof { - return Ok(FeMessage::HealthCheck); - } else { - return Err(err); - } + let mut stream = AsyncPeekable::new(stream); + + if let Err(err) = stream.peek_exact(&mut [0; 1]).await { + // If the stream is empty, it can be a health check. Do not return error. + if err.kind() == ErrorKind::UnexpectedEof { + return Ok(FeMessage::HealthCheck); + } else { + return Err(err); } - }; - assert_eq!(filled1, 1); - - let mut buffer2 = vec![0; 3]; - let filled2 = stream.read_exact(&mut buffer2).await?; - assert_eq!(filled2, 3); - - let mut buffer3 = BytesMut::with_capacity(4); - buffer3.put_slice(&buffer1); - buffer3.put_slice(&buffer2); - - let len = NetworkEndian::read_i32(&buffer3); + } + let len = stream.read_i32().await?; let protocol_num = stream.read_i32().await?; let payload_len = (len - 8) as usize; if payload_len >= isize::MAX as usize {