// 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.

#pragma once

#include <atomic>

#include <chrono>

#include <condition_variable>

#include <cstdint>

#include <functional>

#include <iostream>

#include <mutex>

#include <shared_mutex>

#include <string>

#include <thread>

#include <unordered_map>

#include "common/status.h"

namespace doris {

// Same as

// fe/fe-core/src/main/java/org/apache/doris/common/DNSCache.java

class DNSCache {

public:

    // (hostname, out_ip, is_ipv6, out_gai_err) -> Status.

    // out_gai_err receives the raw getaddrinfo() return code so the cache can tell an

    // authoritative "no such host" (EAI_NONAME) apart from a transient resolver failure.

    using Resolver = std::function<Status(const std::string&, std::string&, bool, int*)>;

    // Per-hostname failure bookkeeping. Only tracked for hostnames currently present in

    // `cache`, which keeps the map bounded (invariant: keys(failure_count) ⊆ keys(cache)).

    struct FailureState {

        // Consecutive resolution failures of any kind. Drives the log throttle.

        uint32_t count = 0;

        // Whether the most recent failure was an authoritative NXDOMAIN. Eviction requires

        // this to be true, so that a DNS-server outage (which yields EAI_AGAIN for every

        // hostname at once) degrades to the stale cached IP instead of wiping the cache.

        bool last_authoritative = false;

    };

    DNSCache();

    // Test-only constructor: uses a custom resolver and does NOT start the

    // background refresh thread.  Call refresh_for_test() to drive one cycle.

    explicit DNSCache(Resolver resolver);

    ~DNSCache();

    // get ip by hostname

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

private:

    // Resolve hostname to IP address.

    // 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 is set to true when DNS returned a live result, false when the

    // returned IP is the stale cached fallback from a failed lookup.

    std::string _resolve_hostname(const std::string& hostname, bool* is_fresh = nullptr);

    // update the ip of hostname in cache; out_state (if non-null) is set to the

    // current failure bookkeeping read under the same lock; out_ip (if non-null)

    // receives the resolved IP so callers can use it without a second cache lookup

    // (avoids operator[] mutation under shared_lock).

    Status _update(const std::string& hostname, FailureState* out_state = nullptr,

                   std::string* out_ip = nullptr);

    // erase a hostname from cache unconditionally (with unique_lock)

    void _erase(const std::string& hostname);

    // Erase a hostname from cache only if it still meets the eviction criteria:

    // failure_count >= threshold AND the most recent failure was authoritative.

    // Re-reads the live state under the same lock that performs the erase, so a

    // concurrent successful resolution that cleared it is not lost.

    // Returns true if the host was erased, false otherwise.

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

    // Drop a hostname from the cache and write a negative-cache tombstone (subject to

    // dns_cache_negative_ttl_seconds). Caller must already hold a unique_lock on `mutex`.

    void _evict_locked(const std::string& hostname);

    // Record a tombstone for a hostname that could not be resolved and has no cached IP,

    // so repeated get() calls do not each pay a full blocking getaddrinfo. Uses

    // try_emplace so an entry just re-armed by get()'s single-flight path is preserved.

    void _remember_unresolvable(const std::string& hostname);

    // one refresh cycle: update every cached hostname and evict if needed

    void _refresh_once();

    // a function for refresh daemon thread

    // update cache at fix internal

    void _refresh_cache();

    // ── test helpers (accessible via friend class DNSCacheTest) ──────────────

    size_t size_for_test() const {

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

        return cache.size();

    }

    size_t negative_cache_size_for_test() const {

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

        return _negative_cache.size();

    }

    uint32_t failure_count_for_test(const std::string& hostname) const {

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

        auto it = failure_count.find(hostname);

        return it != failure_count.end() ? it->second.count : 0;

    }

    // Run one refresh cycle synchronously (no sleep).  Only meaningful when

    // the object was constructed with the test constructor (no background thread).

    void refresh_for_test() { _refresh_once(); }

    // Backdate all negative-cache entries far into the past so they appear

    // expired without removing them.  Use this to simulate TTL expiry in tests

    // that need the re-arm path in get() to trigger.

    // Backdating relative to now() (rather than to steady_clock's epoch, which is

    // typically boot time) keeps this correct on a freshly booted host.

    void _expire_negative_cache_for_test() {

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

        auto backdated = std::chrono::steady_clock::now() - std::chrono::hours(24 * 365);

        for (auto& [k, v] : _negative_cache) {

            v = backdated;

        }

    }

    friend class DNSCacheTest;

private:

    Resolver _resolver; // null → use global hostname_to_ip

    // hostname -> ip

    std::unordered_map<std::string, std::string> cache;

    // hostname -> consecutive resolution failure bookkeeping

    std::unordered_map<std::string, FailureState> failure_count;

    // hostname -> eviction timestamp; effective deadline is computed as

    // eviction_time + dns_cache_negative_ttl_seconds to honor live config changes.

    std::unordered_map<std::string, std::chrono::steady_clock::time_point> _negative_cache;

    mutable std::shared_mutex mutex;

    std::thread refresh_thread;

    // Protects stop_refresh and signals _refresh_cache to wake early on destroy.

    std::mutex _cv_mutex;

    std::condition_variable _cv;

    std::atomic<bool> stop_refresh {false};

};

} // end of namespace doris