//  Copyright (c) Meta Platforms, Inc. and affiliates.

//  This source code is licensed under both the GPLv2 (found in the

//  COPYING file in the root directory) and Apache 2.0 License

//  (found in the LICENSE.Apache file in the root directory).

#pragma once

#include "cache/cache_reservation_manager.h"

#include "rocksdb/secondary_cache.h"

namespace ROCKSDB_NAMESPACE {

class CacheWithSecondaryAdapter : public CacheWrapper {

 public:

  explicit CacheWithSecondaryAdapter(

      std::shared_ptr<Cache> target,

      std::shared_ptr<SecondaryCache> secondary_cache,

      TieredAdmissionPolicy adm_policy = TieredAdmissionPolicy::kAdmPolicyAuto,

      bool distribute_cache_res = false);

  ~CacheWithSecondaryAdapter() override;

  Status Insert(

      const Slice& key, ObjectPtr value, const CacheItemHelper* helper,

      size_t charge, Handle** handle = nullptr,

      Priority priority = Priority::LOW,

      const Slice& compressed_value = Slice(),

      CompressionType type = CompressionType::kNoCompression) override;

  Handle* Lookup(const Slice& key, const CacheItemHelper* helper,

                 CreateContext* create_context,

                 Priority priority = Priority::LOW,

                 Statistics* stats = nullptr) override;

  using Cache::Release;

  bool Release(Handle* handle, bool erase_if_last_ref = false) override;

  ObjectPtr Value(Handle* handle) override;

  void StartAsyncLookup(AsyncLookupHandle& async_handle) override;

  void WaitAll(AsyncLookupHandle* async_handles, size_t count) override;

  std::string GetPrintableOptions() const override;

  const char* Name() const override;

  void SetCapacity(size_t capacity) override;

  Status GetSecondaryCacheCapacity(size_t& size) const override;

  Status GetSecondaryCachePinnedUsage(size_t& size) const override;

  Status UpdateCacheReservationRatio(double ratio);

  Status UpdateAdmissionPolicy(TieredAdmissionPolicy adm_policy);

  Cache* TEST_GetCache() { return target_.get(); }

  SecondaryCache* TEST_GetSecondaryCache() { return secondary_cache_.get(); }

 private:

  static constexpr size_t kReservationChunkSize = 1 << 20;

  bool EvictionHandler(const Slice& key, Handle* handle, bool was_hit);

  void StartAsyncLookupOnMySecondary(AsyncLookupHandle& async_handle);

  Handle* Promote(

      std::unique_ptr<SecondaryCacheResultHandle>&& secondary_handle,

      const Slice& key, const CacheItemHelper* helper, Priority priority,

      Statistics* stats, bool found_dummy_entry, bool kept_in_sec_cache);

  bool ProcessDummyResult(Cache::Handle** handle, bool erase);

  void CleanupCacheObject(ObjectPtr obj, const CacheItemHelper* helper);

  std::shared_ptr<SecondaryCache> secondary_cache_;

  TieredAdmissionPolicy adm_policy_;

  // Whether to proportionally distribute cache memory reservations, i.e

  // placeholder entries with null value and a non-zero charge, across

  // the primary and secondary caches.

  bool distribute_cache_res_;

  // A cache reservation manager to keep track of secondary cache memory

  // usage by reserving equivalent capacity against the primary cache

  std::shared_ptr<ConcurrentCacheReservationManager> pri_cache_res_;

  // Fraction of a cache memory reservation to be assigned to the secondary

  // cache

  double sec_cache_res_ratio_;

  // Mutex for use when managing cache memory reservations. Should not be used

  // for other purposes, as it may risk causing deadlocks.

  mutable port::Mutex cache_res_mutex_;

  // Total memory reserved by placeholder entriesin the cache

  size_t placeholder_usage_;

  // Total placeholoder memory charged to both the primary and secondary

  // caches. Will be <= placeholder_usage_.

  size_t reserved_usage_;

  // Amount of memory reserved in the secondary cache. This should be

  // reserved_usage_ * sec_cache_res_ratio_ in steady state.

  size_t sec_reserved_;

};

}  // namespace ROCKSDB_NAMESPACE