intermediate changes

ref:cde9a383711a11544ce7e107a78147fb96cc4029

1 files changed, 198 insertions, 0 deletions

diff --git a/library/cpp/dns/cache.cpp b/library/cpp/dns/cache.cpp
new file mode 100644
index 00000000000..05c14e82fce
--- /dev/null
+++ b/library/cpp/dns/cache.cpp
@@ -0,0 +1,198 @@
+#include "cache.h"
+
+#include "thread.h"
+
+#include <util/system/tls.h>
+#include <util/system/info.h>
+#include <util/system/rwlock.h>
+#include <util/thread/singleton.h>
+#include <util/generic/singleton.h>
+#include <util/generic/hash.h>
+
+using namespace NDns;
+
+namespace {
+    struct TResolveTask {
+        enum EMethod {
+            Normal,
+            Threaded
+        };
+
+        inline TResolveTask(const TResolveInfo& info, EMethod method)
+            : Info(info)
+            , Method(method)
+        {
+        }
+
+        const TResolveInfo& Info;
+        const EMethod Method;
+    };
+
+    class IDns {
+    public:
+        virtual ~IDns() = default;
+        virtual const TResolvedHost* Resolve(const TResolveTask&) = 0;
+    };
+
+    typedef TAtomicSharedPtr<TResolvedHost> TResolvedHostPtr;
+
+    struct THashResolveInfo {
+        inline size_t operator()(const TResolveInfo& ri) const {
+            return ComputeHash(ri.Host) ^ ri.Port;
+        }
+    };
+
+    struct TCompareResolveInfo {
+        inline bool operator()(const NDns::TResolveInfo& x, const NDns::TResolveInfo& y) const {
+            return x.Host == y.Host && x.Port == y.Port;
+        }
+    };
+
+    class TGlobalCachedDns: public IDns, public TNonCopyable {
+    public:
+        const TResolvedHost* Resolve(const TResolveTask& rt) override {
+            //2. search host in cache
+            {
+                TReadGuard guard(L_);
+
+                TCache::const_iterator it = C_.find(rt.Info);
+
+                if (it != C_.end()) {
+                    return it->second.Get();
+                }
+            }
+
+            TResolvedHostPtr res = ResolveA(rt);
+
+            //update cache
+            {
+                TWriteGuard guard(L_);
+
+                std::pair<TCache::iterator, bool> updateResult = C_.insert(std::make_pair(TResolveInfo(res->Host, rt.Info.Port), res));
+                TResolvedHost* rh = updateResult.first->second.Get();
+
+                if (updateResult.second) {
+                    //fresh resolved host, set cache record id for it
+                    rh->Id = C_.size() - 1;
+                }
+
+                return rh;
+            }
+        }
+
+        void AddAlias(const TString& host, const TString& alias) noexcept {
+            TWriteGuard guard(LA_);
+
+            A_[host] = alias;
+        }
+
+        static inline TGlobalCachedDns* Instance() {
+            return SingletonWithPriority<TGlobalCachedDns, 65530>();
+        }
+
+    private:
+        inline TResolvedHostPtr ResolveA(const TResolveTask& rt) {
+            TString originalHost(rt.Info.Host);
+            TString host(originalHost);
+
+            //3. replace host to alias, if exist
+            if (A_.size()) {
+                TReadGuard guard(LA_);
+                TStringBuf names[] = {"*", host};
+
+                for (const auto& name : names) {
+                    TAliases::const_iterator it = A_.find(name);
+
+                    if (it != A_.end()) {
+                        host = it->second;
+                    }
+                }
+            }
+
+            if (host.length() > 2 && host[0] == '[') {
+                TString unbracedIpV6(host.data() + 1, host.size() - 2);
+                host.swap(unbracedIpV6);
+            }
+
+            TAutoPtr<TNetworkAddress> na;
+
+            //4. getaddrinfo (direct or in separate thread)
+            if (rt.Method == TResolveTask::Normal) {
+                na.Reset(new TNetworkAddress(host, rt.Info.Port));
+            } else if (rt.Method == TResolveTask::Threaded) {
+                na = ThreadedResolve(host, rt.Info.Port);
+            } else {
+                Y_ASSERT(0);
+                throw yexception() << TStringBuf("invalid resolve method");
+            }
+
+            return new TResolvedHost(originalHost, *na);
+        }
+
+        typedef THashMap<TResolveInfo, TResolvedHostPtr, THashResolveInfo, TCompareResolveInfo> TCache;
+        TCache C_;
+        TRWMutex L_;
+        typedef THashMap<TString, TString> TAliases;
+        TAliases A_;
+        TRWMutex LA_;
+    };
+
+    class TCachedDns: public IDns {
+    public:
+        inline TCachedDns(IDns* slave)
+            : S_(slave)
+        {
+        }
+
+        const TResolvedHost* Resolve(const TResolveTask& rt) override {
+            //1. search in local thread cache
+            {
+                TCache::const_iterator it = C_.find(rt.Info);
+
+                if (it != C_.end()) {
+                    return it->second;
+                }
+            }
+
+            const TResolvedHost* res = S_->Resolve(rt);
+
+            C_[TResolveInfo(res->Host, rt.Info.Port)] = res;
+
+            return res;
+        }
+
+    private:
+        typedef THashMap<TResolveInfo, const TResolvedHost*, THashResolveInfo, TCompareResolveInfo> TCache;
+        TCache C_;
+        IDns* S_;
+    };
+
+    struct TThreadedDns: public TCachedDns {
+        inline TThreadedDns()
+            : TCachedDns(TGlobalCachedDns::Instance())
+        {
+        }
+    };
+
+    inline IDns* ThrDns() {
+        return FastTlsSingleton<TThreadedDns>();
+    }
+}
+
+namespace NDns {
+    const TResolvedHost* CachedResolve(const TResolveInfo& ri) {
+        TResolveTask rt(ri, TResolveTask::Normal);
+
+        return ThrDns()->Resolve(rt);
+    }
+
+    const TResolvedHost* CachedThrResolve(const TResolveInfo& ri) {
+        TResolveTask rt(ri, TResolveTask::Threaded);
+
+        return ThrDns()->Resolve(rt);
+    }
+
+    void AddHostAlias(const TString& host, const TString& alias) {
+        TGlobalCachedDns::Instance()->AddAlias(host, alias);
+    }
+}