From 2c41c0357d4b25e81479af487dd66e71b34e99d4 Mon Sep 17 00:00:00 2001
From: Arthur Passos <arthur.ti@outlook.com>
Date: Fri, 21 Oct 2022 14:20:49 -0300
Subject: [PATCH] Make CaresPTRResolver completely thread-safe

---
 src/Common/CaresPTRResolver.h         |  3 +++
 src/Common/DNSPTRResolverProvider.cpp |  5 ++++-
 src/Common/LockedDNSPTRResolver.cpp   | 25 +++++++++++++++++++++++++
 src/Common/LockedDNSPTRResolver.h     | 25 +++++++++++++++++++++++++
 4 files changed, 57 insertions(+), 1 deletion(-)
 create mode 100644 src/Common/LockedDNSPTRResolver.cpp
 create mode 100644 src/Common/LockedDNSPTRResolver.h

diff --git a/src/Common/CaresPTRResolver.h b/src/Common/CaresPTRResolver.h
index 38344d75ade..9df6d7aeb72 100644
--- a/src/Common/CaresPTRResolver.h
+++ b/src/Common/CaresPTRResolver.h
@@ -2,6 +2,7 @@
 
 #include <span>
 #include <poll.h>
+#include <mutex>
 #include "DNSPTRResolver.h"
 
 using ares_channel = struct ares_channeldata *;
@@ -46,6 +47,8 @@ namespace DB
         void process_readable_sockets(std::span<pollfd> readable_sockets);
 
         ares_channel channel;
+
+        static std::mutex mutex;
     };
 }
 
diff --git a/src/Common/DNSPTRResolverProvider.cpp b/src/Common/DNSPTRResolverProvider.cpp
index 97d601a3a78..63d23612fba 100644
--- a/src/Common/DNSPTRResolverProvider.cpp
+++ b/src/Common/DNSPTRResolverProvider.cpp
@@ -1,12 +1,15 @@
 #include "DNSPTRResolverProvider.h"
+#include "LockedDNSPTRResolver.h"
 #include "CaresPTRResolver.h"
 
 namespace DB
 {
     std::shared_ptr<DNSPTRResolver> DNSPTRResolverProvider::get()
     {
-        return std::make_shared<CaresPTRResolver>(
+        static auto resolver = std::make_shared<CaresPTRResolver>(
             CaresPTRResolver::provider_token {}
         );
+
+        return resolver;
     }
 }
diff --git a/src/Common/LockedDNSPTRResolver.cpp b/src/Common/LockedDNSPTRResolver.cpp
new file mode 100644
index 00000000000..947a109860d
--- /dev/null
+++ b/src/Common/LockedDNSPTRResolver.cpp
@@ -0,0 +1,25 @@
+#include "LockedDNSPTRResolver.h"
+
+namespace DB {
+
+    std::mutex LockedPTRResolver::mutex;
+
+    LockedPTRResolver::LockedPTRResolver(std::unique_ptr<DNSPTRResolver> resolver_)
+    : resolver(std::move(resolver_))
+    {}
+
+    std::unordered_set<std::string> LockedPTRResolver::resolve(const std::string & ip)
+    {
+        std::lock_guard guard(mutex);
+
+        return resolver->resolve(ip);
+    }
+
+    std::unordered_set<std::string> LockedPTRResolver::resolve_v6(const std::string & ip)
+    {
+        std::lock_guard guard(mutex);
+
+        return resolver->resolve_v6(ip);
+    }
+
+}
diff --git a/src/Common/LockedDNSPTRResolver.h b/src/Common/LockedDNSPTRResolver.h
new file mode 100644
index 00000000000..82fdd189744
--- /dev/null
+++ b/src/Common/LockedDNSPTRResolver.h
@@ -0,0 +1,25 @@
+#pragma once
+
+#include <mutex>
+
+#include <Common/DNSPTRResolver.h>
+
+namespace DB {
+    class LockedPTRResolver : public DNSPTRResolver
+    {
+    public:
+
+        LockedPTRResolver(std::unique_ptr<DNSPTRResolver> resolver);
+
+        std::unordered_set<std::string> resolve(const std::string & ip) override;
+
+        std::unordered_set<std::string> resolve_v6(const std::string & ip) override;
+
+    private:
+        // this needs to be owned
+        std::unique_ptr<DNSPTRResolver> resolver;
+
+        static std::mutex mutex;
+    };
+}
+