Separate factory for test keeper handlers

programs/server/Server.cpp

@@ -57,6 +57,7 @@
 #include <Disks/registerDisks.h>
 #include <Common/Config/ConfigReloader.h>
 #include <Server/HTTPHandlerFactory.h>
+#include <Server/TestKeeperTCPHandlerFactory.h>
 #include "MetricsTransmitter.h"
 #include <Common/StatusFile.h>
 #include <Server/TCPHandlerFactory.h>
@@ -762,7 +763,7 @@ int Server::main(const std::vector<std::string> & /*args*/)
             socket.setReceiveTimeout(settings.receive_timeout);
             socket.setSendTimeout(settings.send_timeout);
             servers.emplace_back(std::make_unique<Poco::Net::TCPServer>(
-                new TCPHandlerFactory(*this, false, false, true),
+                new TestKeeperTCPHandlerFactory(*this),
                 server_pool,
                 socket,
                 new Poco::Net::TCPServerParams));

src/Server/TCPHandlerFactory.h

@@ -5,7 +5,6 @@
 #include <common/logger_useful.h>
 #include <Server/IServer.h>
 #include <Server/TCPHandler.h>
-#include <Server/TestKeeperTCPHandler.h>
 
 namespace Poco { class Logger; }
 
@@ -18,7 +17,6 @@ private:
     IServer & server;
     bool parse_proxy_protocol = false;
     Poco::Logger * log;
-    bool test_keeper;
 
     class DummyTCPHandler : public Poco::Net::TCPServerConnection
     {
@@ -32,10 +30,9 @@ public:
       * and set the information about forwarded address accordingly.
       * See https://github.com/wolfeidau/proxyv2/blob/master/docs/proxy-protocol.txt
       */
-    TCPHandlerFactory(IServer & server_, bool secure_, bool parse_proxy_protocol_, bool test_keeper_ = false)
+    TCPHandlerFactory(IServer & server_, bool secure_, bool parse_proxy_protocol_)
         : server(server_), parse_proxy_protocol(parse_proxy_protocol_)
         , log(&Poco::Logger::get(std::string("TCP") + (secure_ ? "S" : "") + "HandlerFactory"))
-        , test_keeper(test_keeper_)
     {
     }
 
@@ -45,10 +42,7 @@ public:
         {
             LOG_TRACE(log, "TCP Request. Address: {}", socket.peerAddress().toString());
 
-            if (test_keeper)
-                return new TestKeeperTCPHandler(server, socket);
-            else
-                return new TCPHandler(server, socket, parse_proxy_protocol);
+            return new TCPHandler(server, socket, parse_proxy_protocol);
         }
         catch (const Poco::Net::NetException &)
         {

src/Server/TestKeeperTCPHandlerFactory.h

@@ -0,0 +1,44 @@
+#pragma once
+#include <Server/TestKeeperTCPHandler.h>
+#include <Poco/Net/TCPServerConnectionFactory.h>
+#include <Poco/Net/NetException.h>
+#include <common/logger_useful.h>
+#include <Server/IServer.h>
+
+namespace DB
+{
+
+class TestKeeperTCPHandlerFactory : public Poco::Net::TCPServerConnectionFactory
+{
+private:
+    IServer & server;
+    Poco::Logger * log;
+    class DummyTCPHandler : public Poco::Net::TCPServerConnection
+    {
+    public:
+        using Poco::Net::TCPServerConnection::TCPServerConnection;
+        void run() override {}
+    };
+public:
+    TestKeeperTCPHandlerFactory(IServer & server_)
+        : server(server_)
+        , log(&Poco::Logger::get("TestKeeperTCPHandlerFactory"))
+    {
+    }
+
+    Poco::Net::TCPServerConnection * createConnection(const Poco::Net::StreamSocket & socket) override
+    {
+        try
+        {
+            LOG_TRACE(log, "Test keeper request. Address: {}", socket.peerAddress().toString());
+            return new TestKeeperTCPHandler(server, socket);
+        }
+        catch (const Poco::Net::NetException &)
+        {
+            LOG_TRACE(log, "TCP Request. Client is not connected (most likely RST packet was sent).");
+            return new DummyTCPHandler(socket);
+        }
+    }
+};
+
+}