src/dctv/query_cache.h

// Copyright (C) 2020 The Android Open Source Project
//
// Licensed 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 "dctv.h"

#include <boost/intrusive/list.hpp>
#include <cmath>
#include <utility>

#include "hunk.h"
#include "list.h"
#include "mmap.h"
#include "pyutil.h"
#include "unique_fd.h"
#include "vector.h"

namespace dctv {

struct StringTable;
struct QueryExecution;

struct HunkListAccess {
  using HunkList = boost::intrusive::list<
    Hunk,
    boost::intrusive::constant_time_size<false>,
    boost::intrusive::member_hook<Hunk,
                                  Hunk::HunkListLink,
                                  &Hunk::link>
    >;
};

struct QueryCacheConfig final {
  QueryCacheConfig();
  QueryCacheConfig(pyref args, pyref kwargs);
  unique_pyref tmpdir;
  size_t memory_lwm_nbytes = std::numeric_limits<size_t>::max();
  size_t memory_hwm_nbytes = std::numeric_limits<size_t>::max();
  size_t memory_increment_nbytes = 0;
  size_t disk_max_nbytes = 0;
  size_t mmap_min_nbytes;
  size_t block_size = 4 * 1024 * 1024;
  npy_intp max_freelist = 5;
  npy_intp fill_ratio = 3;
  size_t min_disk_spill_nbytes = 8192;
};

// Manages storage for a query session.
struct QueryCache final : BasePyObject,
                          HasPyCtor,
                          SupportsWeakRefs,
                          HasDict,
                          SupportsGcClear

{
  explicit QueryCache(QueryCacheConfig config);
  ~QueryCache() noexcept;

  int py_traverse(visitproc visit, void* arg) const noexcept;
  int py_clear() noexcept;

  // Clear the cache.
  unique_pyref clear();
  void prune();

  // Keep track of QueryExecution instances using this cache.
  void note_execution_in_progress(QueryExecution* qe);
  void note_execution_done(QueryExecution* qe) noexcept;

  inline const QueryCacheConfig& get_config() const noexcept;

  // Return the number of tracked hunks.
  npy_intp get_nhunks() const noexcept;

  // Return the target block size for query operations.  For various
  // reasons, blocks may have different sizes, but we try to use this
  // one as often as possible.
  inline npy_intp get_block_size() const noexcept;

  // Get the string table used for this query cache.  Immutable.
  inline StringTable* get_st() const noexcept;

  // Get the number of bytes currently used to hunk data.
  inline npy_intp get_memory_nbytes() const noexcept;

  // Make an uninitialized hunk-backed array for use in the
  // query system.
  unique_pyarray make_uninit_array(unique_dtype dtype, npy_intp nelem);

  // Set the score to assign to freshly-created hunks for testing
  // purposes.  Set SCORE_OVERRIDE to NAN to disable the override
  // mechanism.  Illegal to call in optimized builds.
  HunkScore get_new_hunk_score_override();
  void set_new_hunk_score_override(HunkScore score_override);

  std::pair<String, UniqueFd> make_cache_file(std::string_view label);
  UniqueFd open_cache_file(const String& name, int mode);
  void delete_cache_file(const String& name) noexcept;

  // Disable cache trim operations
  struct TrimDisabler final : NoCopy {
    explicit inline TrimDisabler(QueryCache* qc) noexcept;
    inline ~TrimDisabler() noexcept;
   private:
    QueryCache* qc;
  };

  static PyTypeObject pytype;
 private:
  // QueryCache is intimately tied to the hunk mechanism.
  friend struct Hunk;

  struct Freelist {
    size_t nbytes;
    List<MmapShared> spare;
  };

  Freelist* find_fl(size_t sz) noexcept;
  void report() noexcept;
  void trim_memory(size_t target_memory_nbytes,
                   Hunk* protect_hunk /*nullable*/);

  // Note changes to memory use.  Called by Hunk when creating,
  // resizing, deflating, inflating, and destroying hunk data.
  void account_memory(size_t nbytes, Hunk* protect_hunk /*nullable*/);
  void unaccount_memory(size_t nbytes) noexcept;

  // Allocate a memory mapping of at least NBYTES bytes for backing a hunk.
  MmapShared allocate_mmap(size_t nbytes);

  // Determine whether an allocation should use the regular C heap or
  // the mmap-chunk machinery for an allocation of a particular size.
  bool should_use_heap(size_t nbytes) const noexcept;

  // Determine whether we should compress a blob to memory or just
  // spill it to disk.
  bool should_spill_to_memory(size_t nbytes) const noexcept;

  // Resize a memory mapping, preserving contents.  The returned
  // memory mapping will have room for at least NEW_NBYTES bytes, but
  // may also be larger.  We try to avoid copying when possible.
  // Does _not_ do memory accounting; it's the caller's job to call
  // account_memory() or unaccount_memory() as appropriate.
  void resize_mmap(MmapShared* map, size_t new_nbytes);

  // Return a memory mapping allocated using allocate_mmap() to the
  // free list.
  void donate_mmap(MmapShared for_freelist) noexcept;

  QueryCacheConfig config;
  unique_obj_pyref<StringTable> st;

  Vector<Freelist> freelists;
  Vector<QueryExecution*> executions;
  HunkScore new_hunk_score_override = NAN;

  HunkListAccess::HunkList hunks;

  npy_intp memory_nbytes = 0;
  npy_intp disk_nbytes = 0;
  UniqueFd cachedir;
  bool disable_trim = false;
  int64_t file_name_counter = 0;

  static PyGetSetDef pygetset[];
  static PyMemberDef pymembers[];
  static PyMethodDef pymethods[];
};

void init_query_cache(pyref m);

}  // namespace dctv

#include "query_cache-inl.h"