event_data_batch.hpp

// Copyright (c) Microsoft Corporation.
// Licensed under the MIT License.
 
#pragma once
#include "models/event_data.hpp"
 
#include <azure/core/amqp/models/amqp_message.hpp>
#include <azure/core/diagnostics/logger.hpp>
#include <azure/core/internal/diagnostics/log.hpp>
 
#include <mutex>
#include <stdexcept>
 
// cspell: words vbin
 
namespace Azure { namespace Messaging { namespace EventHubs { namespace _detail {
  class EventDataBatchFactory;
}}}} // namespace Azure::Messaging::EventHubs::_detail
 
namespace Azure { namespace Messaging { namespace EventHubs {
 
  struct EventDataBatchOptions final
  {
 
    Azure::Nullable<std::uint64_t> MaxBytes;
 
    std::string PartitionKey;
 
    std::string PartitionId;
  };
 
  class EventDataBatch final {
  private:
    const std::string anyPartitionId = "";
 
    std::mutex m_rwMutex;
    std::string m_partitionId;
    std::string m_partitionKey;
    Azure::Nullable<std::uint64_t> m_maxBytes;
    std::vector<std::vector<uint8_t>> m_marshalledMessages;
    // Annotation properties
    const uint32_t BatchedMessageFormat = 0x80013700;
 
    Azure::Core::Amqp::Models::AmqpMessage m_batchEnvelope;
    size_t m_currentSize;
 
  public:
    EventDataBatch(EventDataBatch const& other)
        // Copy constructor cannot be defaulted because of m_rwMutex.
        : m_rwMutex{}, m_partitionId{other.m_partitionId}, m_partitionKey{other.m_partitionKey},
          m_maxBytes{other.m_maxBytes}, m_marshalledMessages{other.m_marshalledMessages},
          m_batchEnvelope{other.m_batchEnvelope}, m_currentSize(other.m_currentSize){};
 
    EventDataBatch& operator=(EventDataBatch const& other)
    {
      // Assignment operator cannot be defaulted because of m_rwMutex.
      if (this != &other)
      {
        m_partitionId = other.m_partitionId;
        m_partitionKey = other.m_partitionKey;
        m_maxBytes = other.m_maxBytes;
        m_marshalledMessages = other.m_marshalledMessages;
        m_batchEnvelope = other.m_batchEnvelope;
        m_currentSize = other.m_currentSize;
      }
      return *this;
    }
 
    std::string GetPartitionId() const { return m_partitionId; }
 
    std::string GetPartitionKey() const { return m_partitionKey; }
 
    uint64_t GetMaxBytes() const { return m_maxBytes.Value(); }
 
    _azure_NODISCARD bool TryAdd(
        std::shared_ptr<Azure::Core::Amqp::Models::AmqpMessage const> const& message)
    {
      return TryAddAmqpMessage(message);
    }
 
    _azure_NODISCARD bool TryAdd(Azure::Messaging::EventHubs::Models::EventData const& message);
 
    size_t NumberOfEvents()
    {
      std::lock_guard<std::mutex> lock(m_rwMutex);
      return m_marshalledMessages.size();
    }
 
    Azure::Core::Amqp::Models::AmqpMessage ToAmqpMessage() const;
 
  private:
    bool TryAddAmqpMessage(
        std::shared_ptr<Azure::Core::Amqp::Models::AmqpMessage const> const& message);
 
    size_t CalculateActualSizeForPayload(std::vector<uint8_t> const& payload)
    {
      const size_t vbin8Overhead = 5;
      const size_t vbin32Overhead = 8;
 
      if (payload.size() < 256)
      {
        return payload.size() + vbin8Overhead;
      }
      return payload.size() + vbin32Overhead;
    }
 
    Azure::Core::Amqp::Models::AmqpMessage CreateBatchEnvelope(
        std::shared_ptr<Azure::Core::Amqp::Models::AmqpMessage const> const& message) const
    {
      // Create the batch envelope from the prototype message. This copies all the attributes
      // *except* the body attribute to the batch envelope.
      Azure::Core::Amqp::Models::AmqpMessage batchEnvelope{*message};
      batchEnvelope.BodyType = Azure::Core::Amqp::Models::MessageBodyType::None;
      batchEnvelope.MessageFormat = BatchedMessageFormat;
      return batchEnvelope;
    }
 
    EventDataBatch(EventDataBatchOptions options = {})
        : m_partitionId{options.PartitionId}, m_partitionKey{options.PartitionKey},
          m_maxBytes{options.MaxBytes}, m_marshalledMessages{}, m_batchEnvelope{}, m_currentSize{0}
    {
      if (!options.PartitionId.empty() && !options.PartitionKey.empty())
      {
        throw std::runtime_error("Either PartionID or PartitionKey can be set, but not both.");
      }
 
      if (options.PartitionId.empty())
      {
        m_partitionId = anyPartitionId;
      }
    };
 
    friend class _detail::EventDataBatchFactory;
  };
}}} // namespace Azure::Messaging::EventHubs