diff --git a/include/hareflow/detail/accumulator.h b/include/hareflow/detail/accumulator.h index 377f9ce..99510a1 100644 --- a/include/hareflow/detail/accumulator.h +++ b/include/hareflow/detail/accumulator.h @@ -9,6 +9,7 @@ #include #include "hareflow/detail/internal_types.h" +#include "hareflow/detail/semaphore.h" namespace hareflow::detail { @@ -17,12 +18,14 @@ struct AccumulatedMessage { std::chrono::steady_clock::time_point publish_time, MessagePtr message, std::vector encoded_message, - ConfirmationHandler confirmation_handler) + ConfirmationHandler confirmation_handler, + Semaphore::Permit enqueue_permit) : publishing_id(publishing_id), publish_time(publish_time), message(std::move(message)), encoded_message(std::move(encoded_message)), - confirmation_handler(std::move(confirmation_handler)) + confirmation_handler(std::move(confirmation_handler)), + enqueue_permit(std::move(enqueue_permit)) { } @@ -31,6 +34,7 @@ struct AccumulatedMessage { MessagePtr message; const std::vector encoded_message; ConfirmationHandler confirmation_handler; + Semaphore::Permit enqueue_permit; }; class AccumulatorDestroyedException : public std::runtime_error @@ -43,7 +47,7 @@ class Accumulator public: Accumulator(std::uint32_t capacity, CodecPtr codec); - bool add(std::uint64_t publishing_id, MessagePtr message, ConfirmationHandler confirmation_handler); + bool add(std::uint64_t publishing_id, MessagePtr message, ConfirmationHandler confirmation_handler, Semaphore::Permit enqueue_permit); std::vector extract_all(); void set_max_frame_size(std::uint32_t max_frame_size); diff --git a/include/hareflow/detail/semaphore.h b/include/hareflow/detail/semaphore.h index 159ffa7..272918d 100644 --- a/include/hareflow/detail/semaphore.h +++ b/include/hareflow/detail/semaphore.h @@ -5,6 +5,8 @@ #include #include #include +#include +#include namespace hareflow::detail { @@ -16,11 +18,56 @@ class SemaphoreDestroyedException : public std::runtime_error class Semaphore { public: + class Permit + { + public: + Permit() : m_semaphore(nullptr) + { + } + Permit(const Permit&) = delete; + Permit(Permit&& other) : m_semaphore(other.m_semaphore) + { + other.m_semaphore = nullptr; + } + ~Permit() + { + release(); + } + + Permit& operator=(const Permit&) = delete; + Permit& operator=(Permit&& other) + { + if (this != &other) { + release(); + m_semaphore = other.m_semaphore; + other.m_semaphore = nullptr; + } + return *this; + } + + private: + friend class Semaphore; + + Permit(Semaphore& semaphore) : m_semaphore(&semaphore) + { + } + + void release() + { + if (m_semaphore != nullptr) { + m_semaphore->release(); + m_semaphore = nullptr; + } + } + + Semaphore* m_semaphore; + }; + Semaphore(std::uint32_t desired) : m_lock(), m_released(), m_available(desired), m_destroyed(false) { } - void acquire() + Permit acquire() { std::unique_lock lock{m_lock}; m_released.wait(lock, [this] { return m_available > 0 || m_destroyed; }); @@ -28,40 +75,42 @@ class Semaphore throw SemaphoreDestroyedException{"Semaphore was destroyed"}; } --m_available; + return Permit{*this}; } - bool try_acquire_for(std::chrono::milliseconds duration) + std::optional try_acquire_for(std::chrono::milliseconds duration) { std::unique_lock lock{m_lock}; bool acquired = m_released.wait_for(lock, duration, [this] { return m_available > 0 || m_destroyed; }); if (m_destroyed) { throw SemaphoreDestroyedException{"Semaphore was destroyed"}; } - if (acquired) { - --m_available; + if (!acquired) { + return std::nullopt; } - return acquired; + --m_available; + return Permit{*this}; } - void release(std::uint32_t count = 1) + void destroy() { { std::unique_lock lock{m_lock}; - m_available += count; + m_destroyed = true; } m_released.notify_all(); } - void destroy() +private: + void release() { { std::unique_lock lock{m_lock}; - m_destroyed = true; + ++m_available; } m_released.notify_all(); } -private: std::mutex m_lock; std::condition_variable m_released; std::uint32_t m_available; diff --git a/src/accumulator.cpp b/src/accumulator.cpp index 03f59eb..53cae4f 100644 --- a/src/accumulator.cpp +++ b/src/accumulator.cpp @@ -19,7 +19,7 @@ Accumulator::Accumulator(std::uint32_t capacity, CodecPtr codec) m_pending.reserve(m_capacity); } -bool Accumulator::add(std::uint64_t publishing_id, MessagePtr message, ConfirmationHandler confirmation_handler) +bool Accumulator::add(std::uint64_t publishing_id, MessagePtr message, ConfirmationHandler confirmation_handler, Semaphore::Permit enqueue_permit) { std::vector encoded = m_codec->encode(message); @@ -33,7 +33,8 @@ bool Accumulator::add(std::uint64_t publishing_id, MessagePtr message, Confirmat std::chrono::steady_clock::now(), std::move(message), std::move(encoded), - std::move(confirmation_handler)); + std::move(confirmation_handler), + std::move(enqueue_permit)); m_space_available.wait(lock, [this] { return m_pending.size() < m_capacity || m_destroyed; }); if (m_destroyed) { throw AccumulatorDestroyedException{"Accumulator was destroyed"}; diff --git a/src/producer_impl.cpp b/src/producer_impl.cpp index aeaf189..10ce488 100644 --- a/src/producer_impl.cpp +++ b/src/producer_impl.cpp @@ -95,13 +95,16 @@ void ProducerImpl::send(std::uint64_t publishing_id, MessagePtr message, Confirm } try { + Semaphore::Permit enqueue_permit; if (m_enqueue_timeout == std::chrono::milliseconds::zero()) { - m_enqueue_semaphore.acquire(); - } else if (!m_enqueue_semaphore.try_acquire_for(m_enqueue_timeout)) { + enqueue_permit = m_enqueue_semaphore.acquire(); + } else if (auto permit = m_enqueue_semaphore.try_acquire_for(m_enqueue_timeout)) { + enqueue_permit = std::move(*permit); + } else { throw ProducerException{ProducerErrorCode::MessageEnqueuingFailed, "Enqueueing message timeout"}; } - if (m_accumulator->add(publishing_id, std::move(message), std::move(confirmation_handler))) { + if (m_accumulator->add(publishing_id, std::move(message), std::move(confirmation_handler), std::move(enqueue_permit))) { publish_batch(); } } catch (const SemaphoreDestroyedException&) { @@ -244,7 +247,6 @@ void ProducerImpl::handle_publish_confirm(std::uint64_t publishing_id) if (message != nullptr && message->confirmation_handler != nullptr) { message->confirmation_handler(ConfirmationStatus{publishing_id, message->message, true, static_cast(ResponseCode::Ok)}); } - m_enqueue_semaphore.release(); } void ProducerImpl::handle_publish_error(std::uint64_t publishing_id, ResponseCode error_code) @@ -260,7 +262,6 @@ void ProducerImpl::handle_publish_error(std::uint64_t publishing_id, ResponseCod if (message != nullptr && message->confirmation_handler != nullptr) { message->confirmation_handler(ConfirmationStatus{publishing_id, message->message, false, static_cast(error_code)}); } - m_enqueue_semaphore.release(); } void ProducerImpl::handle_metadata_update(std::string_view stream) @@ -406,4 +407,4 @@ void ProducerImpl::reconnect_to_server() } } -} // namespace hareflow::detail \ No newline at end of file +} // namespace hareflow::detail diff --git a/tests/producer.cpp b/tests/producer.cpp index 8d6e92d..67684c9 100644 --- a/tests/producer.cpp +++ b/tests/producer.cpp @@ -81,10 +81,14 @@ TEST_F(ProducerTest, MessageTooBig) { EXPECT_CALL(m_client_mock, max_frame_size()).WillOnce(testing::Return(100)); - hareflow::ProducerPtr producer = producer_builder().build(); + hareflow::ProducerPtr producer = producer_builder().max_unconfirmed(1).enqueue_timeout(std::chrono::seconds{1}).build(); auto message = hareflow::MessageBuilder().body(std::vector(101)).build(); EXPECT_THAT([&]() { producer->send(message, [](auto...) {}); }, testing::Throws()); + + EXPECT_CALL(m_client_mock, publish_encoded(testing::_, testing::SizeIs(1))); + EXPECT_NO_THROW(producer->send(hareflow::MessageBuilder{}.body("hello").build(), [](auto...) {})); + producer->flush(); } TEST_F(ProducerTest, PeriodicFlush) @@ -151,8 +155,9 @@ TEST_F(ProducerTest, EnqueueTimeout) TEST_F(ProducerTest, ConfirmTimeout) { - hareflow::ProducerPtr producer = producer_builder().confirm_timeout(std::chrono::milliseconds{1}).build(); - EXPECT_CALL(m_client_mock, publish_encoded(testing::_, testing::SizeIs(1))); + hareflow::ProducerPtr producer = + producer_builder().confirm_timeout(std::chrono::milliseconds{1}).max_unconfirmed(1).enqueue_timeout(std::chrono::seconds{1}).build(); + EXPECT_CALL(m_client_mock, publish_encoded(testing::_, testing::SizeIs(1))).Times(2); std::promise invoked; producer->send(hareflow::MessageBuilder{}.body("hello").build(), [&](auto& confirmation_status) { @@ -163,6 +168,9 @@ TEST_F(ProducerTest, ConfirmTimeout) producer->flush(); EXPECT_EQ(invoked.get_future().wait_for(std::chrono::seconds{1}), std::future_status::ready); + + EXPECT_NO_THROW(producer->send(hareflow::MessageBuilder{}.body("hello").build(), [](auto...) {})); + producer->flush(); } TEST_F(ProducerTest, StopFailsUnconfirmed) @@ -308,4 +316,4 @@ TEST_F(ProducerTest, StreamDeleted) EXPECT_THAT( [&]() { producer->send(hareflow::MessageBuilder{}.body("hello").build(), [](auto...) {}); }, testing::Throws(testing::Property(&hareflow::ProducerException::code, hareflow::ProducerErrorCode::ProducerStopped))); -} \ No newline at end of file +}