ceph/src/arrow/r/src/memorypool.cpp

   1 // Licensed to the Apache Software Foundation (ASF) under one
   2 // or more contributor license agreements.  See the NOTICE file
   3 // distributed with this work for additional information
   4 // regarding copyright ownership.  The ASF licenses this file
   5 // to you under the Apache License, Version 2.0 (the
   6 // "License"); you may not use this file except in compliance
   7 // with the License.  You may obtain a copy of the License at
   8 //
   9 //   http://www.apache.org/licenses/LICENSE-2.0
  10 //
  11 // Unless required by applicable law or agreed to in writing,
  12 // software distributed under the License is distributed on an
  13 // "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
  14 // KIND, either express or implied.  See the License for the
  15 // specific language governing permissions and limitations
  16 // under the License.
  17
  18 #include "./arrow_types.h"
  19 #if defined(ARROW_R_WITH_ARROW)
  20 #include <arrow/memory_pool.h>
  21 #include <arrow/util/mutex.h>
  22
  23 class GcMemoryPool : public arrow::MemoryPool {
  24  public:
  25   GcMemoryPool() : pool_(arrow::default_memory_pool()) {}
  26
  27   arrow::Status Allocate(int64_t size, uint8_t** out) override {
  28     return GcAndTryAgain([&] { return pool_->Allocate(size, out); });
  29   }
  30
  31   arrow::Status Reallocate(int64_t old_size, int64_t new_size, uint8_t** ptr) override {
  32     return GcAndTryAgain([&] { return pool_->Reallocate(old_size, new_size, ptr); });
  33   }
  34
  35   void Free(uint8_t* buffer, int64_t size) override { pool_->Free(buffer, size); }
  36
  37   int64_t bytes_allocated() const override { return pool_->bytes_allocated(); }
  38
  39   int64_t max_memory() const override { return pool_->max_memory(); }
  40
  41   std::string backend_name() const override { return pool_->backend_name(); }
  42
  43  private:
  44   template <typename Call>
  45   arrow::Status GcAndTryAgain(const Call& call) {
  46     if (call().ok()) {
  47       return arrow::Status::OK();
  48     } else {
  49       auto lock = mutex_.Lock();
  50
  51       // ARROW-10080: Allocation may fail spuriously since the garbage collector is lazy.
  52       // Force it to run then try again in case any reusable allocations have been freed.
  53       static cpp11::function gc = cpp11::package("base")["gc"];
  54       gc();
  55     }
  56     return call();
  57   }
  58
  59   arrow::util::Mutex mutex_;
  60   arrow::MemoryPool* pool_;
  61 };
  62
  63 static GcMemoryPool g_pool;
  64
  65 arrow::MemoryPool* gc_memory_pool() { return &g_pool; }
  66
  67 // [[arrow::export]]
  68 std::shared_ptr<arrow::MemoryPool> MemoryPool__default() {
  69   return std::shared_ptr<arrow::MemoryPool>(&g_pool, [](...) {});
  70 }
  71
  72 // [[arrow::export]]
  73 double MemoryPool__bytes_allocated(const std::shared_ptr<arrow::MemoryPool>& pool) {
  74   return pool->bytes_allocated();
  75 }
  76
  77 // [[arrow::export]]
  78 double MemoryPool__max_memory(const std::shared_ptr<arrow::MemoryPool>& pool) {
  79   return pool->max_memory();
  80 }
  81
  82 // [[arrow::export]]
  83 std::string MemoryPool__backend_name(const std::shared_ptr<arrow::MemoryPool>& pool) {
  84   return pool->backend_name();
  85 }
  86
  87 // [[arrow::export]]
  88 std::vector<std::string> supported_memory_backends() {
  89   return arrow::SupportedMemoryBackendNames();
  90 }
  91
  92 #endif