src/frontend/memory_pool.cc¶

// Copyright 2026 The Tenzr Project Authors
//
// SPDX-License-Identifier: Apache-2.0
//
// 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
//
//      https://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.

#include "tenzr/frontend/memory_pool.h"

#include <sys/mman.h>
#include <unistd.h>

#include <algorithm>
#include <atomic>
#include <cstdlib>

#include "tenzr/common/error.h"

namespace tenzr {

  PinnedBuffer::PinnedBuffer(size_t size_bytes, bool page_locked) : size_bytes_(size_bytes), page_locked_(false) {
    if (size_bytes_ == 0) {
      return;
    }

    size_t page_size = static_cast<size_t>(sysconf(_SC_PAGESIZE));
    if (page_size == 0) {
      page_size = 4096;
    }
    size_t aligned_size = ((size_bytes_ + page_size - 1) / page_size) * page_size;

    int ret = posix_memalign(&data_, page_size, aligned_size);
    TENZR_CHECK(ret == 0 && data_ != nullptr, "posix_memalign failed to allocate pinned buffer of size " + std::to_string(size_bytes_));

    if (page_locked) {
      if (mlock(data_, size_bytes_) == 0) {
        page_locked_ = true;
      }
    }
  }

  PinnedBuffer::~PinnedBuffer() {
    if (deleter_) {
#if TENZR_EXCEPTIONS_ENABLED
      try {
        deleter_();
      } catch (...) {
        // External deleters must not throw; swallow to keep the destructor noexcept.
      }
#else
      deleter_();
#endif
      data_ = nullptr;
      return;
    }
    if (data_) {
      if (page_locked_) {
        munlock(data_, size_bytes_);
      }
      std::free(data_);
      data_ = nullptr;
    }
  }

  std::shared_ptr<PinnedBuffer> PinnedBuffer::Wrap(void* data, size_t size_bytes, std::function<void()> deleter) {
    auto buf = std::shared_ptr<PinnedBuffer>(new PinnedBuffer());
    buf->data_ = data;
    buf->size_bytes_ = size_bytes;
    buf->page_locked_ = false;
    buf->deleter_ = std::move(deleter);
    return buf;
  }

  PinnedBuffer::PinnedBuffer(PinnedBuffer&& other) noexcept
      : data_(other.data_), size_bytes_(other.size_bytes_), page_locked_(other.page_locked_), deleter_(std::move(other.deleter_)) {
    other.data_ = nullptr;
    other.size_bytes_ = 0;
    other.page_locked_ = false;
    other.deleter_ = nullptr;
  }

  PinnedBuffer& PinnedBuffer::operator=(PinnedBuffer&& other) noexcept {
    if (this != &other) {
      if (deleter_) {
        deleter_();
      } else if (data_) {
        if (page_locked_) {
          munlock(data_, size_bytes_);
        }
        std::free(data_);
      }
      data_ = other.data_;
      size_bytes_ = other.size_bytes_;
      page_locked_ = other.page_locked_;
      deleter_ = std::move(other.deleter_);
      other.data_ = nullptr;
      other.size_bytes_ = 0;
      other.page_locked_ = false;
      other.deleter_ = nullptr;
    }
    return *this;
  }

  namespace {
    std::atomic<bool> s_pool_alive{false};
  }

  MemoryPool::MemoryPool() { s_pool_alive.store(true, std::memory_order_release); }

  MemoryPool::~MemoryPool() {
    s_pool_alive.store(false, std::memory_order_release);
    Clear();
  }

  MemoryPool& MemoryPool::Instance() {
    static MemoryPool pool;
    return pool;
  }

  std::shared_ptr<PinnedBuffer> MemoryPool::Allocate(size_t size_bytes, bool page_locked) {
    auto return_to_pool = [size_bytes](PinnedBuffer* ptr) {
      if (!s_pool_alive.load(std::memory_order_acquire)) {
        delete ptr;
        return;
      }
      std::lock_guard<std::mutex> g(MemoryPool::Instance().mutex_);
      if (!s_pool_alive.load(std::memory_order_acquire)) {
        delete ptr;
        return;
      }
      MemoryPool::Instance().free_buffers_[size_bytes].push_back(std::unique_ptr<PinnedBuffer>(ptr));
    };

    std::unique_lock<std::mutex> lock(mutex_);
    auto& bucket = free_buffers_[size_bytes];
    if (!bucket.empty()) {
      auto buf = std::move(bucket.back());
      bucket.pop_back();
      return std::shared_ptr<PinnedBuffer>(buf.release(), return_to_pool);
    }
    lock.unlock();

    auto raw_buf = std::make_unique<PinnedBuffer>(size_bytes, page_locked);
    return std::shared_ptr<PinnedBuffer>(raw_buf.release(), return_to_pool);
  }

  void MemoryPool::Clear() {
    std::lock_guard<std::mutex> lock(mutex_);
    free_buffers_.clear();
  }

}  // namespace tenzr