// 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 <netdb.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) {
                    // No stack trace: this is an expected steady state, not an anomaly, and
                    // it is returned once per caller request for as long as the host stays
                    // unresolvable. Capturing a stack here would make Status::Error() log a
                    // WARNING per call (and every caller logs status.to_string() again),
                    // recreating exactly the be.WARNING flood this cache exists to stop.
                    return Status::InternalError<false>(
                            "Hostname {} is in negative DNS cache (recently evicted or "
                            "unresolvable), 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 {
                    // Lost the single-flight race; see above for why this carries no stack.
                    return Status::InternalError<false>(
                            "Hostname {} is in negative DNS cache (recently evicted or "
                            "unresolvable), 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;
    int gai_err = 0;
    Status status =
            _resolver ? _resolver(hostname, resolved_ip, BackendOptions::is_bind_ipv6(), &gai_err)
                      : hostname_to_ip(hostname, resolved_ip, BackendOptions::is_bind_ipv6(),
                                       &gai_err);

    if (!status.ok() || resolved_ip.empty()) {
        if (is_fresh) {
            *is_fresh = false;
        }
        // EAI_NONAME is the resolver authoritatively answering "this name does not exist",
        // which is the only evidence that a backend is really gone. Everything else
        // (EAI_AGAIN = resolver unreachable or timed out, EAI_SYSTEM, EAI_FAIL, ...) means
        // DNS itself is unhealthy while the host is most likely still up at its last known
        // address, so those failures must never lead to eviction — otherwise a resolver
        // outage would wipe every hostname at once and turn a DNS incident into a
        // cluster-wide RPC outage.
        const bool authoritative = (gai_err == EAI_NONAME);
        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()) {
                    FailureState& state = failure_count[hostname];
                    // The counter tracks failures of any kind so the throttled log below
                    // stays informative during a resolver outage; only `last_authoritative`
                    // gates eviction.
                    failures = ++state.count;
                    state.last_authoritative = authoritative;
                }
            }
            // 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 << ", error: "
                                 << (authoritative ? "NXDOMAIN, host is gone"
                                                   : "transient, DNS unhealthy")
                                 << "), 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::_evict_locked(const std::string& hostname) {
    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();
    }
}

void DNSCache::_remember_unresolvable(const std::string& hostname) {
    int32_t ttl = config::dns_cache_negative_ttl_seconds;
    if (ttl <= 0) {
        return;
    }
    std::unique_lock<std::shared_mutex> lock(mutex);
    // try_emplace, not operator[]: get()'s single-flight path may have just re-armed this
    // tombstone to now(); overwriting it would be harmless there but would also let two
    // callers racing on the same host each reset the deadline, loosening the rate limit.
    _negative_cache.try_emplace(hostname, std::chrono::steady_clock::now());
}

void DNSCache::_erase(const std::string& hostname) {
    std::unique_lock<std::shared_mutex> lock(mutex);
    _evict_locked(hostname);
}

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.count < threshold ||
        !fc_it->second.last_authoritative) {
        // Either a concurrent successful resolution cleared or reset the counter between
        // _update() returning and this call — do not erase a now-healthy entry — or the
        // most recent failure was transient (resolver unreachable) rather than an
        // authoritative NXDOMAIN, in which case the host is probably still alive.
        return false;
    }
    _evict_locked(hostname);
    return true;
}

Status DNSCache::_update(const std::string& hostname, FailureState* out_state,
                         std::string* out_ip) {
    bool is_fresh = false;
    std::string real_ip = _resolve_hostname(hostname, &is_fresh);
    if (real_ip.empty()) {
        if (out_state) {
            *out_state = FailureState {};
        }
        if (out_ip) {
            out_ip->clear();
        }
        // The host could not be resolved and has no cached IP to fall back on, so it never
        // entered `cache` and will therefore never reach the eviction path that writes a
        // tombstone. Record one here as well: otherwise every single get() on a hostname
        // that has never resolved (a typo in the FE, a backend registered before its DNS
        // record propagated) pays a full blocking getaddrinfo, and many of those calls run
        // on bthreads where long blocking is especially costly.
        _remember_unresolvable(hostname);
        return Status::InternalError<false>(
                "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_state) {
            *out_state = FailureState {};
        }
        if (out_ip) {
            out_ip->clear();
        }
        // No stack trace: like the negative-cache hits in get(), this is an expected
        // outcome that can repeat on every request while the host stays evicted.
        return Status::InternalError<false>(
                "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_state) {
        auto fc_it = failure_count.find(hostname);
        *out_state = fc_it != failure_count.end() ? fc_it->second : FailureState {};
    }
    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) {
        // Each _update() below performs a blocking getaddrinfo, so one cycle over a
        // cluster whose DNS is timing out can take minutes. Without this check the
        // destructor's join() would be held up for exactly that long, which would
        // undo the point of making the wait itself interruptible.
        if (stop_refresh.load(std::memory_order_acquire)) {
            break;
        }
        FailureState state;
        Status st = _update(key, &state);
        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`.
        // `last_authoritative` keeps this restricted to hosts the resolver has positively
        // reported as non-existent; a DNS outage yields transient errors for every host at
        // once and must leave the cache (and the stale-IP fallback) intact.
        int32_t threshold = config::dns_cache_max_consecutive_failures;
        if (threshold > 0 && state.last_authoritative &&
            state.count >= 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 "
                             << state.count << " 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