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

#include <cstddef>

#include <cstdint>

#include <deque>

#include <functional>

#include <memory>

#include <mutex>

#include <vector>

#include "common/atomic_shared_ptr.h"

#include "common/status.h"

#include "io/cache/file_cache_common.h"

#include "runtime/memory/mem_tracker_limiter.h"

#include "util/slice.h"

namespace doris {

class ThreadPool;

}

namespace doris::io {

class BlockFileCache;

class AsyncCacheWriteEpochRegistry;

class InflightWriteBufferIndex;

/// Cache admission attributes captured on the query thread and replayed by a write worker.

struct CacheAdmissionContext {

    /// Query identity used by per-query cache admission and accounting.

    TUniqueId query_id;

    FileCacheType cache_type {FileCacheType::NORMAL};

    int64_t expiration_time {0};

    int64_t tablet_id {0};

    bool is_warmup {false};

    /// Capture only the cache-admission fields that remain valid after the query thread returns.

    static CacheAdmissionContext from_cache_context(const CacheContext& context, int64_t tablet_id);

    /// Recreate the worker-local cache context and attach its temporary statistics sink.

    CacheContext to_cache_context(ReadStatistics* stats) const;

};

/// Reference-counted payload whose allocation is charged to the async-write memory tracker.

class AsyncCacheWriteBuffer {

public:

    ~AsyncCacheWriteBuffer();

    char* data() { return _data; }

    const char* data() const { return _data; }

    size_t size() const { return _size; }

private:

    friend class AsyncCacheWriteManager;

    AsyncCacheWriteBuffer(size_t size, std::shared_ptr<MemTrackerLimiter> tracker);

    char* _data = nullptr;

    size_t _size = 0;

    std::shared_ptr<MemTrackerLimiter> _tracker;

};

using AsyncCacheWriteBufferPtr = std::shared_ptr<AsyncCacheWriteBuffer>;

/// One live per-file write generation. All async read plans and writes for the same cache key share

/// the current token. A key-scoped remove invalidates that token and lets later reads capture a new

/// generation without retaining an epoch entry for every historical cache key.

class AsyncCacheWriteEpochToken {

public:

    ~AsyncCacheWriteEpochToken();

    uint64_t generation() const { return _generation; }

    bool is_valid() const { return _valid.load(std::memory_order_acquire); }

private:

    friend class AsyncCacheWriteEpochRegistry;

    AsyncCacheWriteEpochToken(const UInt128Wrapper& cache_hash, uint64_t generation,

                              std::weak_ptr<AsyncCacheWriteEpochRegistry> registry);

    UInt128Wrapper _cache_hash;

    uint64_t _generation {0};

    std::weak_ptr<AsyncCacheWriteEpochRegistry> _registry;

    std::atomic<bool> _valid {true};

};

/// Composite persistence fence shared by an async read plan and its derived write tasks.

/// `cache_epoch` invalidates all older writes on one cache disk during a full-cache clear, while

/// `key_token` lets a single-file removal invalidate only older writes for that cache key. The

/// epoch does not govern inflight reads because a cache hash identifies immutable file content.

struct AsyncCacheWriteEpoch {

    uint64_t cache_epoch {0};

    std::shared_ptr<AsyncCacheWriteEpochToken> key_token;

};

enum class AsyncCacheWriteBlockSubmitResult : uint8_t {

    /// The manager accepted ownership of a complete block buffer.

    SUBMITTED,

    /// The cache invalidation fence changed before submission.

    STALE_EPOCH,

    /// The copy-based submission could not allocate its tracked block buffer.

    BUFFER_ALLOCATION_FAILED,

    /// The same cache block already has a buffer indexed by an accepted write task.

    ALREADY_INFLIGHT,

    /// The manager is unavailable or its admission policy found no queue capacity.

    REJECTED,

};

enum class AsyncCacheWriteAdmissionMode : uint8_t {

    /// Normal work may replace the oldest queued task when pending capacity is full.

    ALLOW_OLDEST_QUEUED_EVICTION,

    /// Lower-priority work is accepted only from currently unused pending capacity.

    REQUIRE_SPARE_CAPACITY,

};

/// Transient source bytes and immutable metadata for one complete cache-block submission.

/// `data` is copied before try_submit_block() returns. `buffer_size` is the fixed cache block

/// allocation size; only the physical EOF block may provide a shorter data slice.

struct AsyncCacheWriteBlockRequest {

    UInt128Wrapper cache_hash;

    size_t file_offset {0};

    Slice data;

    size_t buffer_size {0};

    CacheAdmissionContext admission_ctx;

    AsyncCacheWriteEpoch write_epoch;

    /// Optional index whose lifetime must cover every accepted write task.

    InflightWriteBufferIndex* inflight_index {nullptr};

};

/// Owned complete cache-block payload. The buffer must come from allocate_tracked_buffer() on the

/// target manager and is transferred without another data copy.

struct AsyncCacheWriteOwnedBlockRequest {

    UInt128Wrapper cache_hash;

    size_t file_offset {0};

    size_t write_size {0};

    AsyncCacheWriteBufferPtr buffer;

    CacheAdmissionContext admission_ctx;

    AsyncCacheWriteEpoch write_epoch;

    AsyncCacheWriteAdmissionMode admission_mode {

            AsyncCacheWriteAdmissionMode::ALLOW_OLDEST_QUEUED_EVICTION};

    /// Optional index whose lifetime must cover every accepted write task.

    InflightWriteBufferIndex* inflight_index {nullptr};

};

/// One cache-block write. The sole production submitter allocates every `buffer` with exactly

/// `file_cache_each_block_size` bytes. `write_size` is the valid prefix starting at `file_offset`;

/// only the physical EOF block may use less than the full buffer. `write_epoch` prevents a worker

/// from resurrecting data after cache invalidation.

struct AsyncCacheWriteTask {

    using Finalizer = std::function<void(const AsyncCacheWriteTask&)>;

    UInt128Wrapper cache_hash;

    size_t file_offset {0};

    size_t write_size {0};

    AsyncCacheWriteBufferPtr buffer;

    CacheAdmissionContext admission_ctx;

    int64_t submit_ts_us {0};

    AsyncCacheWriteEpoch write_epoch;

    Finalizer on_finalized;

    /// Assert the task contract at the manager boundary.

    void validate() const;

    /// Return the full allocation capacity charged to pending-byte accounting.

    size_t buffer_size() const;

    /// Run the optional owner cleanup after this task reaches a terminal state.

    void finalize() const;

};

/// Complete per-cache-disk worker and memory settings. The manager receives this value

/// explicitly at construction and through update_options(); it never reads global config.

struct AsyncCacheWriteManagerOptions {

    size_t worker_count {1};

    // Accepted queued+active buffer capacity. With fixed block-size buffers, any remainder smaller

    // than one block is intentionally unusable.

    size_t max_pending_bytes {1};

};

/// Resolve the configured BE-wide pending-byte ownership limit. A positive value is used

/// unchanged; -1 selects max(1 GiB, 1% of `be_mem_limit`). Config validation rejects every other

/// value.

Status resolve_async_file_cache_write_max_pending_bytes(int64_t configured_bytes,

                                                        int64_t be_mem_limit,

                                                        size_t* resolved_bytes);

/// Owns the bounded async-write queue and workers for one BlockFileCache (one cache disk).

///

/// The referenced cache must outlive this manager. Shutdown stops new producers, waits registered

/// producers, and drains all accepted tasks before worker resources are released.

class AsyncCacheWriteManager {

public:

    /// @param cache Non-owning target cache; it must outlive this manager.

    /// @param options Initial worker and pending-memory limits.

    AsyncCacheWriteManager(BlockFileCache* cache, AsyncCacheWriteManagerOptions options);

    ~AsyncCacheWriteManager();

    /// Create the worker pool and schedule the configured long-running workers. Idempotent.

    Status start();

    /// Admit `task` into the memory-bounded FIFO without waiting for disk I/O. Because all tasks

    /// have one fixed cache-block buffer capacity, normal admission to a full queue displaces

    /// exactly one oldest queued task. REQUIRE_SPARE_CAPACITY instead rejects a full-queue

    /// submission so lower-priority work cannot displace normal writes. Active tasks are never

    /// displaced. After a runtime limit decrease, normal submissions continue to replace the

    /// oldest queued task without increasing pending bytes, even while existing pending bytes

    /// exceed the new limit. This call can briefly wait for the queue mutex and finalizes a

    /// displaced task before returning.

    /// @return true if ownership was transferred to the queue; false when workers have not been

    /// started, during shutdown, or on backpressure. A rejected task's finalization callback is

    /// not invoked.

    bool try_submit(AsyncCacheWriteTask task,

                    AsyncCacheWriteAdmissionMode admission_mode =

                            AsyncCacheWriteAdmissionMode::ALLOW_OLDEST_QUEUED_EVICTION);

    /// Copy and submit one complete cache block through the same epoch, inflight-deduplication,

    /// memory-accounting, and queue path used by remote cache-miss reads.

    /// @return ALREADY_INFLIGHT when another accepted task already owns this block; in that case

    /// the request data is not copied.

    AsyncCacheWriteBlockSubmitResult try_submit_block(AsyncCacheWriteBlockRequest request);

    /// Submit an already owned complete cache block without copying its payload. This is the

    /// handoff used after background hole filling completes a tracked block buffer.

    /// The manager retains the buffer after this call only when SUBMITTED is returned.

    AsyncCacheWriteBlockSubmitResult try_submit_owned_block(

            AsyncCacheWriteOwnedBlockRequest request);

    /// Return whether a fixed-size task fits in currently unused pending capacity. The result is a

    /// point-in-time scheduling hint: a concurrent foreground submission may consume the capacity

    /// before the caller hands off its task.

    bool can_accept_without_eviction(size_t buffer_size) const;

    /// Return how many fixed-size tasks fit in the same point-in-time spare capacity. Background

    /// schedulers use this to avoid launching multiple remote reads against one available slot.

    size_t available_slots_without_eviction(size_t buffer_size) const;

    /// Return whether the manager currently accepts submissions. This is a point-in-time hint;

    /// submission methods still make the race-safe final decision.

    bool accepting() const;

    /// Check inflight buffers first, then probe one aligned cache block. Skip writeback if the

    /// block is already cached or being filled. Each lookup releases its lock before the next;

    /// this call must run outside the hole-fill queue mutex because probe may load cache metadata.

    bool should_skip_block_writeback(const UInt128Wrapper& cache_hash, size_t offset, size_t size,

                                     const CacheAdmissionContext& admission_ctx,

                                     InflightWriteBufferIndex* inflight_index) const;

    /// Allocate `size` payload bytes charged to the manager tracker and return them in `buffer`.

    Status allocate_tracked_buffer(size_t size, AsyncCacheWriteBufferPtr* buffer);

    /// Capture the disk-wide epoch and current live generation for `cache_hash`.

    AsyncCacheWriteEpoch current_write_epoch(const UInt128Wrapper& cache_hash);

    /// Return the current disk-wide epoch used by cache-clear operations.

    uint64_t current_cache_epoch() const { return _cache_epoch.load(std::memory_order_acquire); }

    /// Test whether both levels of `epoch` still accept writes.

    bool is_current_write_epoch(const AsyncCacheWriteEpoch& epoch) const;

    /// Test the epoch and record one stale-drop metric when it is no longer current.

    bool check_write_epoch(const AsyncCacheWriteEpoch& epoch);

    /// Invalidate only queued/inflight work captured for `cache_hash` before this call.

    void invalidate_pending_writes(const UInt128Wrapper& cache_hash);

    /// Advance the disk-wide epoch so every previously captured write becomes stale.

    /// @return The newly active disk-wide epoch.

    uint64_t invalidate_all_pending_writes();

    /// Return cache keys whose current valid generation is retained by at least one plan or task.

    /// Invalidated generations are excluded even while stale tasks finish releasing them.

    size_t active_write_epoch_key_count() const;

    /// Resize the number of active workers. A shrink waits only for retiring worker loops.

    /// @param worker_count Positive target worker count for this cache disk.

    Status resize_workers(size_t worker_count);

    /// Replace all mutable manager settings with one coherent snapshot. Configuration adapters

    /// call this method explicitly; the manager itself has no dependency on global config.

    /// @param options Complete validated settings, including the desired worker count.

    /// @return OK after the new snapshot is active; InvalidArgument for invalid limits, or a

    /// worker-resize error when the requested concurrency cannot be applied.

    Status update_options(const AsyncCacheWriteManagerOptions& options);

    /// Return the currently active settings as a value snapshot.

    AsyncCacheWriteManagerOptions options() const;

    /// Stop submissions, drain all accepted tasks, and join worker loops. Idempotent.

    void shutdown();

    /// Return accepted tasks that have not yet completed finalization.

    size_t pending_count() const { return _pending_count.load(std::memory_order_relaxed); }

    /// Return buffer-capacity bytes owned by queued and active tasks.

    size_t pending_bytes() const { return _pending_bytes.load(std::memory_order_relaxed); }

    /// Return accepted tasks still waiting in the FIFO queue, excluding active workers.

    size_t queued_count() const;

    /// Return buffer-capacity bytes still waiting in the FIFO queue.

    size_t queued_bytes() const { return _queued_bytes.load(std::memory_order_relaxed); }

    /// Return tasks currently owned by workers.

    size_t active_task_count() const { return _active_task_count.load(std::memory_order_relaxed); }

    /// Return buffer-capacity bytes currently owned by workers.

    size_t active_bytes() const { return _active_bytes.load(std::memory_order_relaxed); }

    /// Return worker loops that are currently alive.

    size_t running_worker_count() const {

        return _running_worker_count.load(std::memory_order_relaxed);

    }

    /// Return bytes currently held by tracked task buffers.

    int64_t buffer_memory_bytes() const { return _mem_tracker->consumption(); }

    /// Return tasks displaced by full-queue admission.

    uint64_t evicted_oldest_count() const;

    /// Return the current rolling P99 wait to acquire the FIFO mutex.

    int64_t queue_lock_wait_p99_us() const;

    /// Return the current rolling P99 FIFO mutex critical-section duration.

    int64_t queue_lock_hold_p99_us() const;

private:

    class Worker;

    enum class TaskFinalizationReason : uint8_t {

        WORKER_FINISHED,

        EVICTED_OLDEST,

    };

    /// Owns bvar registration and translates manager events into coherent metric updates.

    class Metrics;

    /// Resize the owned worker set while `_lifecycle_mutex` is held and `_worker_pool` exists.

    /// The pool's minimum thread count is kept equal to the long-running Worker task count so a

    /// task is never accepted without a backing OS thread.

    Status _resize_workers_locked(size_t worker_count);

    /// Stop and join workers in `[keep_worker_count, _workers.size())` while the lifecycle mutex is

    /// held. Stop requests are published under `_queue_mutex` before waking the worker loops.

    void _stop_workers_locked(size_t keep_worker_count);

    /// Process one task already moved from queued to active ownership.

    void _process_task(AsyncCacheWriteTask task);

    /// Move the oldest queued task to active ownership.

    bool _try_activate_task(AsyncCacheWriteTask* task);

    /// Revalidate epoch/cache state and persist the task's still-empty complete blocks.

    Status _persist_task(const AsyncCacheWriteTask& task);

    /// Move one active task to its terminal state outside the queue lock.

    void _complete_active_task(AsyncCacheWriteTask task);

    /// Record the terminal reason and let the task run its owner cleanup without the queue lock.

    void _complete_task(AsyncCacheWriteTask task, TaskFinalizationReason reason);

    BlockFileCache* _cache;

    atomic_shared_ptr<const AsyncCacheWriteManagerOptions> _options;

    std::deque<AsyncCacheWriteTask> _queue;

    mutable std::mutex _queue_mutex;

    std::condition_variable _queue_cv;

    // Learned from the first submission and protected by `_queue_mutex`.

    size_t _task_buffer_size {0};

    // Pending covers all accepted tasks, while queued and active are its disjoint ownership states:

    //   pending_count = queue.size() + active_task_count

    //   pending_bytes = queued_bytes + active_bytes

    // Byte state is maintained directly and is authoritative for memory admission. Production

    // tasks currently have a fixed cache-block buffer capacity, including a partial EOF write, but

    // byte accounting deliberately does not depend on deriving bytes from task counts.

    std::atomic<size_t> _pending_count {0};

    std::atomic<size_t> _pending_bytes {0};

    std::atomic<size_t> _queued_bytes {0};

    std::atomic<size_t> _active_task_count {0};

    std::atomic<size_t> _active_bytes {0};

    std::atomic<size_t> _running_worker_count {0};

    std::atomic<size_t> _active_get_or_set_count {0};

    std::atomic<size_t> _active_append_count {0};

    std::atomic<size_t> _active_finalize_count {0};

    std::atomic<bool> _accepting {true};

    std::atomic<size_t> _active_submitters {0};

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

    std::atomic<uint64_t> _cache_epoch {1};

    std::shared_ptr<AsyncCacheWriteEpochRegistry> _write_epoch_registry;

    std::shared_ptr<MemTrackerLimiter> _mem_tracker;

    std::unique_ptr<Metrics> _metrics;

    std::unique_ptr<ThreadPool> _worker_pool;

    std::atomic<size_t> _configured_worker_count {0};

    // Serializes start, resize, and shutdown, including all changes to `_workers`.

    std::mutex _lifecycle_mutex;

    // Protected by `_lifecycle_mutex`. Worker stop state is owned by each Worker.

    std::vector<std::shared_ptr<Worker>> _workers;

};

} // namespace doris::io