AimRT/_deps/libunifex-src/test/submit_allocator_customisation_test.cpp
2025-01-12 20:42:42 +08:00

115 lines
3.1 KiB
C++

/*
* Copyright (c) Facebook, Inc. and its affiliates.
*
* Licensed under the Apache License Version 2.0 with LLVM Exceptions
* (the "License"); you may not use this file except in compliance with
* the License. You may obtain a copy of the License at
*
* https://llvm.org/LICENSE.txt
*
* 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 <unifex/just.hpp>
#include <unifex/scheduler_concepts.hpp>
#include <unifex/single_thread_context.hpp>
#include <unifex/submit.hpp>
#include <unifex/sync_wait.hpp>
#include <unifex/then.hpp>
#include <unifex/via.hpp>
#include <unifex/when_all.hpp>
#include <unifex/with_allocator.hpp>
#include <unifex/memory_resource.hpp>
#include <memory>
#include <gtest/gtest.h>
using namespace unifex;
namespace {
#if !UNIFEX_NO_MEMORY_RESOURCE
using namespace unifex::pmr;
class counting_memory_resource : public memory_resource {
public:
explicit counting_memory_resource(memory_resource *r) noexcept : inner_(r) {}
std::size_t total_allocated_bytes() const { return allocated_.load(); }
std::size_t total_allocation_count() const { return count_.load(); }
private:
void *do_allocate(std::size_t bytes, std::size_t alignment) override {
void *p = inner_->allocate(bytes, alignment);
allocated_ += bytes;
++count_;
return p;
}
void do_deallocate(void *p, std::size_t bytes,
std::size_t alignment) override {
allocated_ -= bytes;
inner_->deallocate(p, bytes, alignment);
}
bool do_is_equal(const memory_resource &other) const noexcept override {
return &other == this;
}
memory_resource *inner_;
std::atomic<std::size_t> allocated_ = 0;
std::atomic<std::size_t> count_ = 0;
};
#endif
template <typename Scheduler, typename Allocator>
void test(Scheduler scheduler, Allocator allocator) {
int value = 0;
auto addToValue = [&](int x) {
// The via() is expected to allocate when it calls submit().
// NOTE: This may start failing if we ever merge via() and typed_via().
return via(scheduler, just(x))
| then([&](int x) {
std::printf("got %i\n", x);
value += x;
});
};
when_all(addToValue(1), addToValue(2))
| with_allocator(allocator)
| sync_wait();
EXPECT_EQ(value, 3);
}
} // anonymous namespace
TEST(with_allocator, SubmitWithStdAllocator) {
single_thread_context thread;
test(thread.get_scheduler(), std::allocator<std::byte>{});
}
#if !UNIFEX_NO_MEMORY_RESOURCE
TEST(with_allocator, SubmitWithCountingAllocator) {
counting_memory_resource res{new_delete_resource()};
polymorphic_allocator<char> alloc{&res};
{
single_thread_context thread;
test(thread.get_scheduler(), alloc);
}
// Check that it freed all the memory it allocated
EXPECT_EQ(res.total_allocated_bytes(), 0);
// Check that it actually called allocate()
EXPECT_EQ(res.total_allocation_count(), 2);
}
#endif