sharable_binary_header.h

Go to the documentation of this file.

/*
 * Copyright 2024 MusicScience37 (Kenta Kabashima)
 *
 * 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 <atomic>
#include <cassert>
#include <cstddef>
#include <cstdlib>
#include <limits>
#include <new>
 
#include "msgpack_rpc/messages/impl/sharable_binary_header_fwd.h"  // IWYU pragma: keep
 
namespace msgpack_rpc::messages::impl {

struct SharableBinaryHeader {
public:
    std::size_t total_memory_size;

    std::size_t binary_size;

    std::size_t binary_capacity;

    alignas(alignof(std::atomic<int>))  // NOLINTNEXTLINE
        char reference_count_buffer[sizeof(std::atomic<int>)];

    bool is_reference_count_enabled;
};

constexpr std::size_t SHARABLE_BINARY_ALIGNMENT = alignof(SharableBinaryHeader);
 
// Assumption to use std::malloc, std::realloc, std::free.
static_assert(SHARABLE_BINARY_ALIGNMENT <= alignof(std::max_align_t));

[[nodiscard]] inline SharableBinaryHeader* allocate_sharable_binary(
    std::size_t binary_size) {
    constexpr std::size_t header_size = sizeof(SharableBinaryHeader);
    const std::size_t total_memory_size = header_size + binary_size;
 
    void* raw_ptr = std::malloc(total_memory_size);  // NOLINT(*-no-malloc)
    if (raw_ptr == nullptr) {
        // Since memory allocation has failed, the class of exceptions with
        // memory allocation like MsgpackRPCException cannot be used.
        throw std::bad_alloc();
    }
 
    auto* ptr = static_cast<SharableBinaryHeader*>(raw_ptr);
    ptr->total_memory_size = total_memory_size;
    ptr->binary_size = binary_size;
    ptr->binary_capacity = binary_size;
    ptr->is_reference_count_enabled = false;
 
    return ptr;
}

[[nodiscard]] inline SharableBinaryHeader* allocate_shared_binary() {
    constexpr std::size_t header_size = sizeof(SharableBinaryHeader);
    // This value is not tuned.
    constexpr std::size_t total_memory_size = 128;
    static_assert(total_memory_size > header_size);
 
    void* raw_ptr = std::malloc(total_memory_size);  // NOLINT(*-no-malloc)
    if (raw_ptr == nullptr) {
        // Since memory allocation has failed, the class of exceptions with
        // memory allocation like MsgpackRPCException cannot be used.
        throw std::bad_alloc();
    }
 
    auto* ptr = static_cast<SharableBinaryHeader*>(raw_ptr);
    ptr->total_memory_size = total_memory_size;
    ptr->binary_size = 0U;
    ptr->binary_capacity = total_memory_size - header_size;
    ptr->is_reference_count_enabled = false;
 
    return ptr;
}

[[nodiscard]] inline SharableBinaryHeader* expand_sharable_binary(
    SharableBinaryHeader* buffer, std::size_t required_size) {
    // std::realloc is not safe when reference count (use of std::atomic<int>)
    // is enabled.
    assert(!buffer->is_reference_count_enabled);
 
    std::size_t new_capacity = buffer->binary_capacity;
    while (true) {
        // TODO Prevent overflow.
        new_capacity <<= 1U;
        if (new_capacity >= required_size) {
            break;
        }
    }
    constexpr std::size_t header_size = sizeof(SharableBinaryHeader);
    const std::size_t new_total_memory_size = header_size + new_capacity;
 
    void* raw_ptr =
        std::realloc(buffer, new_total_memory_size);  // NOLINT(*-no-malloc)
    if (raw_ptr == nullptr) {
        // Since memory allocation has failed, the class of exceptions with
        // memory allocation like MsgpackRPCException cannot be used.
        throw std::bad_alloc();
    }
 
    auto* ptr = static_cast<SharableBinaryHeader*>(raw_ptr);
    ptr->total_memory_size = new_total_memory_size;
    ptr->binary_capacity = new_total_memory_size - header_size;
 
    return ptr;
}

[[nodiscard]] inline std::atomic<int>& reference_count_of(
    SharableBinaryHeader* buffer) noexcept {
    assert(buffer->is_reference_count_enabled);
    // NOLINTNEXTLINE(*-reinterpret-cast)
    return *reinterpret_cast<std::atomic<int>*>(buffer->reference_count_buffer);
}

inline void deallocate_sharable_binary(SharableBinaryHeader* buffer) noexcept {
    if (buffer == nullptr) {
        return;
    }
    if (buffer->is_reference_count_enabled) {
        using AtomicType = std::atomic<int>;
        reference_count_of(buffer).~AtomicType();
    }
    std::free(buffer);  // NOLINT(*-no-malloc)
}

inline void enable_reference_count_of_sharable_binary(
    SharableBinaryHeader* buffer) noexcept {
    assert(!buffer->is_reference_count_enabled);
    new (static_cast<void*>(buffer->reference_count_buffer))
        std::atomic<int>(1);
    buffer->is_reference_count_enabled = true;
}

inline void add_reference_count_of_sharable_binary(
    SharableBinaryHeader* buffer) noexcept {
    reference_count_of(buffer).fetch_add(1, std::memory_order_relaxed);
}

inline void decrease_reference_count_of_sharable_binary(
    SharableBinaryHeader* buffer) noexcept {
    const int reference_count_before =
        reference_count_of(buffer).fetch_sub(1, std::memory_order_acq_rel);
    if (reference_count_before == 1) {
        deallocate_sharable_binary(buffer);
    }
}

inline char* binary_buffer_of(SharableBinaryHeader* buffer) noexcept {
    // NOLINTNEXTLINE
    return reinterpret_cast<char*>(buffer) + sizeof(SharableBinaryHeader);
}
 
}  // namespace msgpack_rpc::messages::impl