Return SessionPool sessions after JSON errors

src/client/SessionPool.cpp

@@ -77,18 +77,24 @@ std::string SessionPool::executeJsonWithParameter(
     const std::string &stmt, const std::unordered_map<std::string, Value> &parameters) {
   auto result = getIdleSession();
   if (result.second) {
-    auto resp = result.first.executeJsonWithParameter(stmt, parameters);
-    auto obj = folly::parseJson(resp);
-    auto respSpaceName = obj["results"][0]["spaceName"].asString();
-    if (respSpaceName != config_.spaceName_) {
-      if (respSpaceName == "") {
-        // This can be caused by reconnect, switch to origin space and try to query again
-        resp = result.first.executeJsonWithParameter("USE " + config_.spaceName_ + ";" + stmt,
-                                                     parameters);
-      } else {
-        // switch to origin space only
-        result.first.execute("USE " + config_.spaceName_);
+    std::string resp;
+    try {
+      resp = result.first.executeJsonWithParameter(stmt, parameters);
+      auto obj = folly::parseJson(resp);
+      auto respSpaceName = obj["results"][0]["spaceName"].asString();
+      if (respSpaceName != config_.spaceName_) {
+        if (respSpaceName == "") {
+          // This can be caused by reconnect, switch to origin space and try to query again
+          resp = result.first.executeJsonWithParameter("USE " + config_.spaceName_ + ";" + stmt,
+                                                       parameters);
+        } else {
+          // switch to origin space only
+          result.first.execute("USE " + config_.spaceName_);
+        }
       }
+    } catch (...) {
+      giveBack(std::move(result.first));
+      throw;
     }
     giveBack(std::move(result.first));
     return resp;

src/client/tests/SessionPoolTest.cpp

@@ -3,17 +3,36 @@
 // This source code is licensed under Apache 2.0 License.
 
 #include <common/Init.h>
+#include <common/datatypes/DataSet.h>
 #include <common/datatypes/Geography.h>
 #include <folly/json.h>
 #include <folly/synchronization/Baton.h>
 #include <glog/logging.h>
 #include <gtest/gtest.h>
 #include <gtest/gtest_prod.h>
-#include <nebula/client/SessionPool.h>
+#include <nebula/client/Config.h>
+#include <nebula/client/Connection.h>
 
 #include <atomic>
+#include <cassert>
 #include <chrono>
+#include <cstddef>
+#include <exception>
+#include <list>
+#include <mutex>
 #include <thread>
+#include <utility>
+#include <vector>
+
+#define private public
+#include <nebula/client/Session.h>
+#undef private
+
+#include <nebula/client/ConnectionPool.h>
+
+#define private public
+#include <nebula/client/SessionPool.h>
+#undef private
 
 #include "./ClientTest.h"
 
@@ -154,6 +173,28 @@ TEST_F(SessionPoolTest, Normal) {
   }
 }
 
+TEST_F(SessionPoolTest, JsonParseFailureReturnsSessionToPool) {
+  nebula::SessionPoolConfig config;
+  config.username_ = "root";
+  config.password_ = "nebula";
+  config.addrs_ = {kServerHost ":9669"};
+  config.spaceName_ = "session_pool_test";
+  config.maxSize_ = 1;
+
+  nebula::SessionPool pool(config);
+  ASSERT_TRUE(pool.init());
+  ASSERT_EQ(pool.idleSessions_.size(), 1);
+
+  // Force Connection::executeJsonWithParameter() to return an empty response.
+  pool.idleSessions_.front().conn_.close();
+  EXPECT_THROW(pool.executeJson("YIELD 1"), std::exception);
+  ASSERT_EQ(pool.idleSessions_.size(), 1);
+
+  auto resp = pool.execute("YIELD 1");
+  EXPECT_NE(resp.errorMsg ? *resp.errorMsg : "", "No idle session.");
+  EXPECT_EQ(pool.idleSessions_.size(), 1);
+}
+
 // case 3
 TEST_F(SessionPoolTest, SwitchSpace) {
   nebula::SessionPoolConfig config;