// Licensed to the Apache Software Foundation (ASF) under one

// or more contributor license agreements.  See the NOTICE file

// distributed with this work for additional information

// regarding copyright ownership.  The ASF licenses this file

// to you under the Apache License, Version 2.0 (the

// "License"); you may not use this file except in compliance

// with the License.  You may obtain a copy of the License at

//

//   http://www.apache.org/licenses/LICENSE-2.0

//

// Unless required by applicable law or agreed to in writing,

// software distributed under the License is distributed on an

// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY

// KIND, either express or implied.  See the License for the

// specific language governing permissions and limitations

// under the License.

#include "util/dns_cache.h"

#include <algorithm>

#include <atomic>

#include <unordered_set>

#include "common/config.h"

#include "service/backend_options.h"

#include "util/network_util.h"

namespace doris {

DNSCache::DNSCache() {

    refresh_thread = std::thread(&DNSCache::_refresh_cache, this);

}

DNSCache::DNSCache(Resolver resolver) : _resolver(std::move(resolver)) {}

DNSCache::~DNSCache() {

    {

        std::lock_guard<std::mutex> lk(_cv_mutex);

        stop_refresh = true;

    }

    _cv.notify_all();

    if (refresh_thread.joinable()) {

        refresh_thread.join();

    }

}

Status DNSCache::get(const std::string& hostname, std::string* ip) {

    bool has_negative_entry = false;

    {

        std::shared_lock<std::shared_mutex> lock(mutex);

        auto it = cache.find(hostname);

        if (it != cache.end()) {

            *ip = it->second;

            return Status::OK();

        }

        auto neg_it = _negative_cache.find(hostname);

        if (neg_it != _negative_cache.end()) {

            int32_t ttl = config::dns_cache_negative_ttl_seconds;

            if (ttl > 0) {

                auto deadline = neg_it->second + std::chrono::seconds(ttl);

                if (std::chrono::steady_clock::now() < deadline) {

                    return Status::InternalError(

                            "Hostname {} is in negative DNS cache (recently evicted), "

                            "skipping resolve",

                            hostname);

                }

            }

            has_negative_entry = true;

        }

    }

    // If the host was in the negative cache with an expired (or disabled) TTL,

    // claim the single-flight retry under unique_lock before the blocking DNS

    // call.  Re-arming the eviction_time to now() makes concurrent callers see

    // an unexpired entry, bounding retries to one per host per TTL period.

    if (has_negative_entry) {

        std::unique_lock<std::shared_mutex> lock(mutex);

        auto neg_it = _negative_cache.find(hostname);

        if (neg_it != _negative_cache.end()) {

            int32_t ttl = config::dns_cache_negative_ttl_seconds;

            if (ttl <= 0) {

                _negative_cache.erase(neg_it);

            } else {

                auto deadline = neg_it->second + std::chrono::seconds(ttl);

                if (std::chrono::steady_clock::now() >= deadline) {

                    neg_it->second = std::chrono::steady_clock::now();

                } else {

                    return Status::InternalError(

                            "Hostname {} is in negative DNS cache (recently evicted), "

                            "skipping resolve",

                            hostname);

                }

            }

        }

    }

    // First access (or negative TTL expired): resolve and populate the cache.

    // Consume the IP returned by _update() directly to avoid a second cache

    // lookup — operator[] under a shared_lock would mutate the map and could

    // reinsert an empty entry if a concurrent refresh cycle evicted the hostname

    // between _update() and here.

    return _update(hostname, nullptr, ip);

}

// Resolve hostname to IP address, similar to Java's DNSCache.resolveHostname.

// If resolution fails, falls back to cached IP if available.

// Returns the resolved IP, or cached IP on failure, or empty string if no cache available.

// *is_fresh (if non-null) is set to true when DNS returned a live result, false

// when the IP comes from the stale cached fallback path.

std::string DNSCache::_resolve_hostname(const std::string& hostname, bool* is_fresh) {

    // Get cached IP first (if any)

    std::string cached_ip;

    {

        std::shared_lock<std::shared_mutex> lock(mutex);

        auto it = cache.find(hostname);

        if (it != cache.end()) {

            cached_ip = it->second;

        }

    }

    // Try to resolve hostname

    std::string resolved_ip;

    Status status = _resolver

                            ? _resolver(hostname, resolved_ip, BackendOptions::is_bind_ipv6())

                            : hostname_to_ip(hostname, resolved_ip, BackendOptions::is_bind_ipv6());

    if (!status.ok() || resolved_ip.empty()) {

        if (is_fresh) *is_fresh = false;

        if (!cached_ip.empty()) {

            // Only track failure counts for hosts that are currently in the cache.

            // Hosts that were never cached or have already been evicted are not

            // tracked, which prevents unbounded growth of failure_count.

            uint32_t failures = 0;

            {

                std::unique_lock<std::shared_mutex> lock(mutex);

                // Re-check that the host is still cached under the unique_lock:

                // it may have been evicted by the refresh thread between our

                // earlier shared_lock read of cached_ip and now (hostname_to_ip

                // can block for seconds on DNS timeout, widening the window).

                // Skipping the bump here preserves keys(failure_count) ⊆ keys(cache).

                if (cache.find(hostname) != cache.end()) {

                    failures = ++failure_count[hostname];

                }

            }

            // Throttle the log: only every N failures or the first failure.

            if (failures > 0) {

                int32_t every_n = std::max(1, config::dns_cache_log_every_n_failures);

                if (failures == 1 || failures % static_cast<uint32_t>(every_n) == 0) {

                    LOG(WARNING) << "Failed to resolve hostname " << hostname

                                 << " (consecutive failures: " << failures

                                 << "), use cached ip: " << cached_ip;

                }

            }

            return cached_ip;

        } else {

            // Throttle to avoid flooding be.WARNING when callers repeatedly

            // query an evicted or never-resolvable hostname.  This branch

            // deliberately does not maintain a per-hostname counter (that

            // would break the keys(failure_count) ⊆ keys(cache) invariant),

            // so the throttle is a coarse global rate limit shared across

            // all hostnames hitting this code path.

            static std::atomic<uint64_t> no_cache_warn_counter {0};

            uint64_t n = no_cache_warn_counter.fetch_add(1, std::memory_order_relaxed) + 1;

            int32_t every_n = std::max(1, config::dns_cache_log_every_n_failures);

            if (n == 1 || n % static_cast<uint64_t>(every_n) == 0) {

                LOG(WARNING) << "Failed to resolve hostname " << hostname

                             << ", no cached ip available";

            }

            return "";

        }

    }

    // Resolution succeeded - clear failure counter for this hostname.

    if (is_fresh) *is_fresh = true;

    {

        std::unique_lock<std::shared_mutex> lock(mutex);

        failure_count.erase(hostname);

    }

    return resolved_ip;

}

void DNSCache::_erase(const std::string& hostname) {

    std::unique_lock<std::shared_mutex> lock(mutex);

    cache.erase(hostname);

    failure_count.erase(hostname);

    int32_t ttl = config::dns_cache_negative_ttl_seconds;

    if (ttl > 0) {

        _negative_cache[hostname] = std::chrono::steady_clock::now();

    }

}

bool DNSCache::_erase_if_still_failing(const std::string& hostname, uint32_t threshold) {

    std::unique_lock<std::shared_mutex> lock(mutex);

    auto fc_it = failure_count.find(hostname);

    if (fc_it == failure_count.end() || fc_it->second < threshold) {

        // A concurrent successful resolution cleared or reset the counter between

        // _update() returning and this call — do not erase a now-healthy entry.

        return false;

    }

    cache.erase(hostname);

    failure_count.erase(hostname);

    int32_t ttl = config::dns_cache_negative_ttl_seconds;

    if (ttl > 0) {

        _negative_cache[hostname] = std::chrono::steady_clock::now();

    }

    return true;

}

Status DNSCache::_update(const std::string& hostname, uint32_t* out_failures, std::string* out_ip) {

    bool is_fresh = false;

    std::string real_ip = _resolve_hostname(hostname, &is_fresh);

    if (real_ip.empty()) {

        if (out_failures) *out_failures = 0;

        if (out_ip) out_ip->clear();

        return Status::InternalError("Failed to resolve hostname {} and no cached ip available",

                                     hostname);

    }

    std::unique_lock<std::shared_mutex> lock(mutex);

    // _resolve_hostname may have captured a stale cached_ip before a concurrent

    // eviction completed.  If the host is now in the negative cache we must not

    // reinsert the stale IP: that would silently undo the eviction and clear the

    // tombstone, defeating the whole purpose of eviction.  Only a fresh DNS

    // result (is_fresh == true, meaning DNS actually resolved) may override an

    // eviction — which indicates the backend is genuinely back.

    if (!is_fresh && _negative_cache.count(hostname)) {

        if (out_failures) *out_failures = 0;

        if (out_ip) out_ip->clear();

        return Status::InternalError(

                "Hostname {} was concurrently evicted; stale-fallback not reinserted", hostname);

    }

    auto it = cache.find(hostname);

    if (it == cache.end() || it->second != real_ip) {

        cache[hostname] = real_ip;

        LOG(INFO) << "update hostname " << hostname << "'s ip to " << real_ip;

    }

    // DNS resolved successfully — remove any negative cache tombstone so

    // subsequent get() calls go straight to the main cache.

    _negative_cache.erase(hostname);

    if (out_ip) *out_ip = real_ip;

    // Read failure_count under the same lock we already hold, so _refresh_once

    // does not need a second lock acquisition to decide on eviction.

    if (out_failures) {

        auto fc_it = failure_count.find(hostname);

        *out_failures = fc_it != failure_count.end() ? fc_it->second : 0;

    }

    return Status::OK();

}

void DNSCache::_refresh_once() {

    std::unordered_set<std::string> keys;

    {

        std::shared_lock<std::shared_mutex> lock(mutex);

        std::transform(cache.begin(), cache.end(), std::inserter(keys, keys.end()),

                       [](const auto& pair) { return pair.first; });

    }

    for (auto& key : keys) {

        uint32_t failures = 0;

        Status st = _update(key, &failures);

        if (!st.ok()) {

            // _update returns an error either when _resolve_hostname returns ""

            // (no fallback IP) or when a stale fallback was suppressed because

            // the host was concurrently evicted.  Either way, log and move on;

            // the threshold check below handles the normal eviction path.

            LOG(WARNING) << "Failed to update DNS cache for hostname " << key << ": "

                         << st.to_string();

        }

        // Evict hostnames that have failed to resolve for too long.

        // This avoids two pathological symptoms after a backend is dropped

        // from the cluster and its DNS record is removed:

        //   1) be.WARNING gets flooded with `failed to get ip from host`.

        //   2) brpc keeps re-using the stale IP from cache, producing

        //      `Fail to wait EPOLLOUT ... Connection timed out`.

        int32_t threshold = config::dns_cache_max_consecutive_failures;

        if (threshold > 0 && failures >= static_cast<uint32_t>(threshold)) {

            // Re-read failure_count under the mutex that also performs the erase

            // to fence any concurrent success that cleared the counter between

            // _update() returning and this point.

            if (_erase_if_still_failing(key, static_cast<uint32_t>(threshold))) {

                LOG(WARNING) << "Evicting hostname " << key << " from DNS cache after " << failures

                             << " consecutive resolution failures";

            }

        }

    }

}

void DNSCache::_refresh_cache() {

    while (!stop_refresh) {

        {

            std::unique_lock<std::mutex> lk(_cv_mutex);

            // Wake up either after 1 minute or when the destructor signals stop.

            _cv.wait_for(lk, std::chrono::minutes(1), [this] { return stop_refresh.load(); });

        }

        if (!stop_refresh) {

            _refresh_once();

        }

    }

}

} // end of namespace doris