From 5bdb5cc02383c6de514b6a71bcd9ea4dc40ce885 Mon Sep 17 00:00:00 2001
From: chenpingzeng <chenpingzeng@huawei.com>
Date: Sat, 28 Nov 2020 10:18:52 +0800
Subject: [PATCH] 1. enable DBserver alive check through ping function 2. check
 availability of DBserver after setup of backend connection

Signed-off-by: chenpingzeng <chenpingzeng@huawei.com>
---
 .../java/io/mycat/net/FrontendConnection.java |  3 +-
 .../net/handler/FrontendAuthenticator.java    |  5 ++
 .../net/handler/FrontendCommandHandler.java   |  4 +-
 .../io/mycat/server/ServerConnection.java     | 60 +++++++++++++++++--
 4 files changed, 65 insertions(+), 7 deletions(-)

diff --git a/src/main/java/io/mycat/net/FrontendConnection.java b/src/main/java/io/mycat/net/FrontendConnection.java
index eb12db7..60b01c5 100644
--- a/src/main/java/io/mycat/net/FrontendConnection.java
+++ b/src/main/java/io/mycat/net/FrontendConnection.java
@@ -459,8 +459,9 @@ public void stmtClose(byte[] data) {
 		}
 	}
 
-	public void ping() {
+	public boolean ping(boolean sendResponse) {
 		write(writeToBuffer(OkPacket.OK, allocate()));
+		return true;
 	}
 
 	public void heartbeat(byte[] data) {
diff --git a/src/main/java/io/mycat/net/handler/FrontendAuthenticator.java b/src/main/java/io/mycat/net/handler/FrontendAuthenticator.java
index e670e6b..bde29e7 100644
--- a/src/main/java/io/mycat/net/handler/FrontendAuthenticator.java
+++ b/src/main/java/io/mycat/net/handler/FrontendAuthenticator.java
@@ -137,6 +137,11 @@ public void handle(byte[] data) {
             return;
         }
 
+        // ping DBserver to make sure backend connection setup successfully
+        if (!source.ping(false)) {
+            failure(ErrorCode.ER_SERVER_SHUTDOWN, "DBserver shutdown.");
+            return;
+        }
         success(auth);
 
     }
diff --git a/src/main/java/io/mycat/net/handler/FrontendCommandHandler.java b/src/main/java/io/mycat/net/handler/FrontendCommandHandler.java
index 305b09d..54fa3c1 100644
--- a/src/main/java/io/mycat/net/handler/FrontendCommandHandler.java
+++ b/src/main/java/io/mycat/net/handler/FrontendCommandHandler.java
@@ -77,7 +77,7 @@ public void handle(byte[] data)
                 break;
             case MySQLPacket.COM_PING:
                 commands.doPing();
-                source.ping();
+                source.ping(true);
                 break;
             case MySQLPacket.COM_QUIT:
                 commands.doQuit();
@@ -119,7 +119,7 @@ public void handle(byte[] data)
                 MycatConfig config = MycatServer.getInstance().getConfig();
                 if( config.getSystem().getIgnoreUnknownCommand()==1){
                     LOGGER.warn("Unknown command:{}",data[4]);
-                    source.ping();
+                    source.ping(true);
                 }else {
                     LOGGER.error("Unknown command:{}",new String(data));
                     source.writeErrMessage(ErrorCode.ER_UNKNOWN_COM_ERROR,
diff --git a/src/main/java/io/mycat/server/ServerConnection.java b/src/main/java/io/mycat/server/ServerConnection.java
index 1fe7b44..812934d 100644
--- a/src/main/java/io/mycat/server/ServerConnection.java
+++ b/src/main/java/io/mycat/server/ServerConnection.java
@@ -26,11 +26,15 @@
 
 import java.io.IOException;
 import java.nio.channels.NetworkChannel;
+import java.sql.Connection;
+import java.sql.ResultSet;
 import java.sql.SQLException;
+import java.sql.Statement;
 import java.util.regex.Matcher;
 import java.util.regex.Pattern;
 
 import io.mycat.backend.jdbc.JDBCConnection;
+import io.mycat.net.mysql.OkPacket;
 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;
 
@@ -167,10 +171,58 @@ public void setLocked(boolean isLocked) {
 		this.isLocked = isLocked;
 	}
 
-	@Override
-	public void ping() {
-		Ping.response(this);
-	}
+    @Override
+    public boolean ping(boolean sendResponse) {
+        if (!(backConn instanceof JDBCConnection)) {
+            if (sendResponse) {
+                write(writeToBuffer(OkPacket.OK, allocate()));
+            }
+            return true;
+        }
+
+        Connection con = ((JDBCConnection)backConn).getCon();
+        ResultSet rsOutput = null;
+        Statement stmtOutput = null;
+        String executeSql = "/* ping */ select 1";
+        boolean isActive = true;
+
+        try {
+            con.setCatalog(catalog);
+            con.setSchema(schema);
+
+            stmtOutput =  con.createStatement();
+            rsOutput   = stmtOutput.executeQuery(executeSql);
+
+            if (sendResponse) {
+                // ping server by select ok, response to client
+                write(writeToBuffer(OkPacket.OK, allocate()));
+            }
+        } catch (SQLException e) {
+            if (sendResponse) {
+                writeErrMessage(ErrorCode.ER_SERVER_SHUTDOWN, "DBServer shutdown.");
+            }
+            
+            isActive = false;
+        } finally {
+            if (rsOutput != null) {
+                try {
+                    rsOutput.close();
+                } catch (SQLException e) {
+                    ; // do nothing
+                }
+            }
+
+            if (stmtOutput != null) {
+                try {
+                    stmtOutput.close();
+                } catch (SQLException e) {
+                    ; // do nothing
+                }
+            }
+        }
+
+        return isActive;
+    }
 
 	@Override
 	public void heartbeat(byte[] data) {