feat: handle Ctrl-C command in MySQL client (#6320)

* feat/answer-ctrl-c-in-mysql:
 ## Implement Connection ID-based Query Killing

 ### Key Changes:
 - **Connection ID Management:**
   - Added `connection_id` to `Session` and `QueryContext` in `src/session/src/lib.rs` and `src/session/src/context.rs`.
   - Updated `MysqlInstanceShim` and `MysqlServer` to handle `connection_id` in `src/servers/src/mysql/handler.rs` and `src/servers/src/mysql/server.rs`.

 - **KILL Statement Enhancements:**
   - Introduced `Kill` enum to handle both `ProcessId` and `ConnectionId` in `src/sql/src/statements/kill.rs`.
   - Updated `ParserContext` to parse `KILL QUERY <connection_id>` in `src/sql/src/parser.rs`.
   - Modified `StatementExecutor` to support killing queries by `connection_id` in `src/operator/src/statement/kill.rs`.

 - **Process Management:**
   - Refactored `ProcessManager` to include `connection_id` in `src/catalog/src/process_manager.rs`.
   - Added `kill_local_process` method for local query termination.

 - **Testing:**
   - Added tests for `KILL` statement parsing and execution in `src/sql/src/parser.rs`.

 ### Affected Files:
 - `Cargo.lock`, `Cargo.toml`
 - `src/catalog/src/process_manager.rs`
 - `src/frontend/src/instance.rs`
 - `src/frontend/src/stream_wrapper.rs`
 - `src/operator/src/statement.rs`
 - `src/operator/src/statement/kill.rs`
 - `src/servers/src/mysql/federated.rs`
 - `src/servers/src/mysql/handler.rs`
 - `src/servers/src/mysql/server.rs`
 - `src/servers/src/postgres.rs`
 - `src/session/src/context.rs`
 - `src/session/src/lib.rs`
 - `src/sql/src/parser.rs`
 - `src/sql/src/statements.rs`
 - `src/sql/src/statements/kill.rs`
 - `src/sql/src/statements/statement.rs`

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>


 Conflicts:

	Cargo.lock

	Cargo.toml

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 ### Enhance Process Management and Execution

 - **`process_manager.rs`**: Added a new method `find_processes_by_connection_id` to filter processes by connection ID, improving process management capabilities.
 - **`kill.rs`**: Refactored the process killing logic to utilize the new `find_processes_by_connection_id` method, streamlining the execution flow and reducing redundant checks.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 ## Commit Message

 ### Update Process ID Type and Refactor Code

 - **Change Process ID Type**: Updated the process ID type from `u64` to `u32` across multiple files to optimize memory usage. Affected files include `process_manager.rs`, `lib.rs`, `database.rs`, `instance.rs`, `server.rs`, `stream_wrapper.rs`, `kill.rs`, `federated.rs`, `handler.rs`, `server.rs`,
 `postgres.rs`, `mysql_server_test.rs`, `context.rs`, `lib.rs`, and `test_util.rs`.

 - **Remove Connection ID**: Removed the `connection_id` field and related logic from `process_manager.rs`, `lib.rs`, `instance.rs`, `server.rs`, `stream_wrapper.rs`, `kill.rs`, `federated.rs`, `handler.rs`, `server.rs`, `postgres.rs`, `mysql_server_test.rs`, `context.rs`, `lib.rs`, and `test_util.rs` to
 simplify the codebase.

 - **Refactor Process Management**: Refactored process management logic to improve clarity and maintainability in `process_manager.rs`, `kill.rs`, and `handler.rs`.

 - **Enhance MySQL Server Handling**: Improved MySQL server handling by integrating process management in `server.rs` and `mysql_server_test.rs`.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 ### Add Process Manager to Postgres Server

 - **`src/frontend/src/server.rs`**: Updated server initialization to include `process_manager`.
 - **`src/servers/src/postgres.rs`**: Modified `MakePostgresServerHandler` to accept `process_id` for session creation.
 - **`src/servers/src/postgres/server.rs`**: Integrated `process_manager` into `PostgresServer` for generating `process_id` during connection handling.
 - **`src/servers/tests/postgres/mod.rs`** and **`tests-integration/src/test_util.rs`**: Adjusted test server setup to accommodate optional `process_manager`.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

* feat/answer-ctrl-c-in-mysql:
 Update `greptime-proto` Dependency

 - Updated the `greptime-proto` dependency to a new revision in both `Cargo.lock` and `Cargo.toml`.
   - `Cargo.lock`: Changed source revision from `d75a56e05a87594fe31ad5c48525e9b2124149ba` to `fdcbe5f1c7c467634c90a1fd1a00a784b92a4e80`.
   - `Cargo.toml`: Updated the `greptime-proto` git revision to match the new commit.

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

---------

Signed-off-by: Lei, HUANG <mrsatangel@gmail.com>

src/servers/src/mysql/handler.rs

@@ -82,6 +82,7 @@ pub struct MysqlInstanceShim {
     user_provider: Option<UserProviderRef>,
     prepared_stmts: Arc<RwLock<HashMap<String, SqlPlan>>>,
     prepared_stmts_counter: AtomicU32,
+    process_id: u32,
 }
 
 impl MysqlInstanceShim {
@@ -89,6 +90,7 @@ impl MysqlInstanceShim {
         query_handler: ServerSqlQueryHandlerRef,
         user_provider: Option<UserProviderRef>,
         client_addr: SocketAddr,
+        process_id: u32,
     ) -> MysqlInstanceShim {
         // init a random salt
         let mut bs = vec![0u8; 20];
@@ -110,12 +112,12 @@ impl MysqlInstanceShim {
                 Some(client_addr),
                 Channel::Mysql,
                 Default::default(),
-                // TODO(sunng87): generate process id properly
-                0,
+                process_id,
             )),
             user_provider,
             prepared_stmts: Default::default(),
             prepared_stmts_counter: AtomicU32::new(1),
+            process_id,
         }
     }
 
@@ -341,6 +343,10 @@ impl<W: AsyncWrite + Send + Sync + Unpin> AsyncMysqlShim<W> for MysqlInstanceShi
         std::env::var("GREPTIMEDB_MYSQL_SERVER_VERSION").unwrap_or_else(|_| "8.4.2".to_string())
     }
 
+    fn connect_id(&self) -> u32 {
+        self.process_id
+    }
+
     fn default_auth_plugin(&self) -> &str {
         self.auth_plugin()
     }

src/servers/src/mysql/server.rs

@@ -18,6 +18,7 @@ use std::sync::Arc;
 
 use async_trait::async_trait;
 use auth::UserProviderRef;
+use catalog::process_manager::ProcessManagerRef;
 use common_runtime::runtime::RuntimeTrait;
 use common_runtime::Runtime;
 use common_telemetry::{debug, warn};
@@ -112,6 +113,7 @@ pub struct MysqlServer {
     spawn_ref: Arc<MysqlSpawnRef>,
     spawn_config: Arc<MysqlSpawnConfig>,
     bind_addr: Option<SocketAddr>,
+    process_manager: Option<ProcessManagerRef>,
 }
 
 impl MysqlServer {
@@ -119,16 +121,23 @@ impl MysqlServer {
         io_runtime: Runtime,
         spawn_ref: Arc<MysqlSpawnRef>,
         spawn_config: Arc<MysqlSpawnConfig>,
+        process_manager: Option<ProcessManagerRef>,
     ) -> Box<dyn Server> {
         Box::new(MysqlServer {
             base_server: BaseTcpServer::create_server("MySQL", io_runtime),
             spawn_ref,
             spawn_config,
             bind_addr: None,
+            process_manager,
         })
     }
 
-    fn accept(&self, io_runtime: Runtime, stream: AbortableStream) -> impl Future<Output = ()> {
+    fn accept(
+        &self,
+        io_runtime: Runtime,
+        stream: AbortableStream,
+        process_manager: Option<ProcessManagerRef>,
+    ) -> impl Future<Output = ()> {
         let spawn_ref = self.spawn_ref.clone();
         let spawn_config = self.spawn_config.clone();
 
@@ -136,7 +145,7 @@ impl MysqlServer {
             let spawn_ref = spawn_ref.clone();
             let spawn_config = spawn_config.clone();
             let io_runtime = io_runtime.clone();
-
+            let process_id = process_manager.as_ref().map(|p| p.next_id()).unwrap_or(8);
             async move {
                 match tcp_stream {
                     Err(e) => warn!(e; "Broken pipe"), // IoError doesn't impl ErrorExt.
@@ -146,7 +155,7 @@ impl MysqlServer {
                         }
                         io_runtime.spawn(async move {
                             if let Err(error) =
-                                Self::handle(io_stream, spawn_ref, spawn_config).await
+                                Self::handle(io_stream, spawn_ref, spawn_config, process_id).await
                             {
                                 warn!(error; "Unexpected error when handling TcpStream");
                             };
@@ -161,10 +170,11 @@ impl MysqlServer {
         stream: TcpStream,
         spawn_ref: Arc<MysqlSpawnRef>,
         spawn_config: Arc<MysqlSpawnConfig>,
+        process_id: u32,
     ) -> Result<()> {
         debug!("MySQL connection coming from: {}", stream.peer_addr()?);
         crate::metrics::METRIC_MYSQL_CONNECTIONS.inc();
-        if let Err(e) = Self::do_handle(stream, spawn_ref, spawn_config).await {
+        if let Err(e) = Self::do_handle(stream, spawn_ref, spawn_config, process_id).await {
             if let Error::InternalIo { error } = &e
                 && error.kind() == std::io::ErrorKind::ConnectionAborted
             {
@@ -184,11 +194,13 @@ impl MysqlServer {
         stream: TcpStream,
         spawn_ref: Arc<MysqlSpawnRef>,
         spawn_config: Arc<MysqlSpawnConfig>,
+        process_id: u32,
     ) -> Result<()> {
         let mut shim = MysqlInstanceShim::create(
             spawn_ref.query_handler(),
             spawn_ref.user_provider(),
             stream.peer_addr()?,
+            process_id,
         );
         let (mut r, w) = stream.into_split();
         let mut w = BufWriter::with_capacity(DEFAULT_RESULT_SET_WRITE_BUFFER_SIZE, w);
@@ -230,7 +242,11 @@ impl Server for MysqlServer {
             .await?;
         let io_runtime = self.base_server.io_runtime();
 
-        let join_handle = common_runtime::spawn_global(self.accept(io_runtime, stream));
+        let join_handle = common_runtime::spawn_global(self.accept(
+            io_runtime,
+            stream,
+            self.process_manager.clone(),
+        ));
         self.base_server.start_with(join_handle).await?;
 
         self.bind_addr = Some(addr);

src/servers/src/postgres.rs

@@ -119,9 +119,13 @@ impl PgWireServerHandlers for PostgresServerHandler {
 }
 
 impl MakePostgresServerHandler {
-    fn make(&self, addr: Option<SocketAddr>) -> PostgresServerHandler {
-        // TODO(sunng87): generate pid from process manager
-        let session = Arc::new(Session::new(addr, Channel::Postgres, Default::default(), 0));
+    fn make(&self, addr: Option<SocketAddr>, process_id: u32) -> PostgresServerHandler {
+        let session = Arc::new(Session::new(
+            addr,
+            Channel::Postgres,
+            Default::default(),
+            process_id,
+        ));
         let handler = PostgresServerHandlerInner {
             query_handler: self.query_handler.clone(),
             login_verifier: PgLoginVerifier::new(self.user_provider.clone()),

src/servers/src/postgres/server.rs

@@ -18,6 +18,7 @@ use std::sync::Arc;
 
 use ::auth::UserProviderRef;
 use async_trait::async_trait;
+use catalog::process_manager::ProcessManagerRef;
 use common_runtime::runtime::RuntimeTrait;
 use common_runtime::Runtime;
 use common_telemetry::{debug, warn};
@@ -37,6 +38,7 @@ pub struct PostgresServer {
     tls_server_config: Arc<ReloadableTlsServerConfig>,
     keep_alive_secs: u64,
     bind_addr: Option<SocketAddr>,
+    process_manager: Option<ProcessManagerRef>,
 }
 
 impl PostgresServer {
@@ -48,6 +50,7 @@ impl PostgresServer {
         keep_alive_secs: u64,
         io_runtime: Runtime,
         user_provider: Option<UserProviderRef>,
+        process_manager: Option<ProcessManagerRef>,
     ) -> PostgresServer {
         let make_handler = Arc::new(
             MakePostgresServerHandlerBuilder::default()
@@ -63,6 +66,7 @@ impl PostgresServer {
             tls_server_config,
             keep_alive_secs,
             bind_addr: None,
+            process_manager,
         }
     }
 
@@ -73,12 +77,12 @@ impl PostgresServer {
     ) -> impl Future<Output = ()> {
         let handler_maker = self.make_handler.clone();
         let tls_server_config = self.tls_server_config.clone();
+        let process_manager = self.process_manager.clone();
         accepting_stream.for_each(move |tcp_stream| {
             let io_runtime = io_runtime.clone();
-
             let tls_acceptor = tls_server_config.get_server_config().map(TlsAcceptor::from);
-
             let handler_maker = handler_maker.clone();
+            let process_id = process_manager.as_ref().map(|p| p.next_id()).unwrap_or(0);
 
             async move {
                 match tcp_stream {
@@ -97,7 +101,7 @@ impl PostgresServer {
 
                         let _handle = io_runtime.spawn(async move {
                             crate::metrics::METRIC_POSTGRES_CONNECTIONS.inc();
-                            let pg_handler = Arc::new(handler_maker.make(addr));
+                            let pg_handler = Arc::new(handler_maker.make(addr, process_id));
                             let r =
                                 process_socket(io_stream, tls_acceptor.clone(), pg_handler).await;
                             crate::metrics::METRIC_POSTGRES_CONNECTIONS.dec();

src/servers/tests/mysql/mysql_server_test.rs

@@ -73,6 +73,7 @@ fn create_mysql_server(table: TableRef, opts: MysqlOpts<'_>) -> Result<Box<dyn S
             0,
             opts.reject_no_database,
         )),
+        None,
     ))
 }
 

src/servers/tests/postgres/mod.rs

@@ -71,6 +71,7 @@ fn create_postgres_server(
         0,
         io_runtime,
         user_provider,
+        None,
     )))
 }