// The default producer configuration only use number of messages to limit batching

return ProducerConfiguration()

.setBatchingType(ProducerConfiguration::KeyBasedBatching)

.setBatchingMaxAllowedSizeInBytes(static_cast<unsigned long>(-1))

.setBatchingMaxPublishDelayMs(3600 * 1000);

protected:

KeyBasedBatchingTest() : client_("pulsar://localhost:6650") {}

void TearDown() override { client_.close(); }

void initTopicName(const std::string& testName) {

topicName_ = "KeyBasedBatchingTest-" + testName + "-" + std::to_string(time(nullptr));

}

void initProducer(const ProducerConfiguration& producerConfig) {

ASSERT_EQ(ResultOk, client_.createProducer(topicName_, producerConfig, producer_));

}

void initConsumer() { ASSERT_EQ(ResultOk, client_.subscribe(topicName_, "SubscriptionName", consumer_)); }

void receiveAndAck(Message& msg) {

ASSERT_EQ(ResultOk, consumer_.receive(msg, 3000));

ASSERT_EQ(ResultOk, consumer_.acknowledge(msg));

}

Client client_;

Producer producer_;

Consumer consumer_;

std::string topicName_;

initTopicName("Flush");

// no limits for batching

initProducer(createDefaultProducerConfig().setBatchingMaxMessages(

static_cast<unsigned int>(-1)) // no limits for batching

);

constexpr int numMessages = 100;

const std::string keys[] = {"A", "B"};

std::atomic_int numMessageSent{0};

for (int i = 0; i < numMessages; i++) {

producer_.sendAsync(MessageBuilder().setOrderingKey(keys[i % 2]).setContent("x").build(),

[&numMessageSent](Result result, const MessageId&) {

numMessageSent++;

ASSERT_EQ(result, ResultOk);

});

}

ASSERT_EQ(ResultOk, producer_.flush());

ASSERT_EQ(numMessageSent.load(), numMessages);

initTopicName("OrderingKeyPriority");

initProducer(createDefaultProducerConfig().setBatchingMaxMessages(3));

initConsumer();

Latch latch(3);

auto sendCallback = [&latch](Result result, const MessageId& id) {

ASSERT_EQ(result, ResultOk);

latch.countdown();

};

// "0" is send to batch of "A" because ordering key has higher priority

producer_.sendAsync(MessageBuilder().setContent("0").setOrderingKey("A").setPartitionKey("B").build(),

sendCallback);

producer_.sendAsync(MessageBuilder().setContent("1").setOrderingKey("A").build(), sendCallback);

producer_.sendAsync(MessageBuilder().setContent("2").setOrderingKey("B").build(), sendCallback);

latch.countdown();

Message msg;

receiveAndAck(msg);

ASSERT_EQ("0", msg.getDataAsString());

ASSERT_EQ("A", msg.getOrderingKey());

ASSERT_EQ("B", msg.getPartitionKey());

receiveAndAck(msg);

ASSERT_EQ("1", msg.getDataAsString());

ASSERT_EQ("A", msg.getOrderingKey());

receiveAndAck(msg);

ASSERT_EQ("2", msg.getDataAsString());

ASSERT_EQ("B", msg.getOrderingKey());

initTopicName("SequenceId");

initProducer(createDefaultProducerConfig().setBatchingMaxMessages(6));

initConsumer();

Latch latch(6);

auto sendAsync = [this, &latch](const std::string& key, const std::string& value) {

producer_.sendAsync(MessageBuilder().setOrderingKey(key).setContent(value).build(),

[&latch](Result result, const MessageId& id) {

ASSERT_EQ(result, ResultOk);

latch.countdown();

});

};

sendAsync("A", "0");

sendAsync("B", "1");

sendAsync("C", "2");

sendAsync("B", "3");

sendAsync("C", "4");

sendAsync("A", "5");

// Batches are sent in ascending order of the first message's sequence id (BatchMessageKeyBasedContainer

// sorts by sendArgs->sequenceId). Send order gives A=0, B=1, C=2 for first per key, so batches: A, B, C.

// A: 0, 5

// B: 1, 3

// C: 2, 4

latch.wait();

std::vector<std::string> receivedKeys;

std::vector<std::string> receivedValues;

for (int i = 0; i < 6; i++) {

Message msg;

receiveAndAck(msg);

receivedKeys.emplace_back(msg.getOrderingKey());

receivedValues.emplace_back(msg.getDataAsString());

}

decltype(receivedKeys) expectedKeys{"A", "A", "B", "B", "C", "C"};

decltype(receivedValues) expectedValues{"0", "5", "1", "3", "2", "4"};

EXPECT_EQ(receivedKeys, expectedKeys);

EXPECT_EQ(receivedValues, expectedValues);

initTopicName("SingleBatch");

initProducer(createDefaultProducerConfig().setBatchingMaxMessages(5));

initConsumer();

constexpr int numMessages = 5 * 100;

std::atomic_int numMessageSent{0};

// messages with no key are packed to the same batch and this batch has no key

// the broker uses `NON_KEY` as the key when dispatching messages from this batch

for (int i = 0; i < numMessages; i++) {

producer_.sendAsync(MessageBuilder().setContent("x").build(),

[&numMessageSent](Result result, const MessageId&) {

ASSERT_EQ(result, ResultOk);

++numMessageSent;

});

}

Message msg;

for (int i = 0; i < numMessages; i++) {

receiveAndAck(msg);

}

ASSERT_EQ(ResultTimeout, consumer_.receive(msg, 3000));

ASSERT_EQ(numMessageSent.load(), numMessages);

initTopicName("CloseBeforeSend");

// Any asynchronous send won't be completed unless `close()` or `flush()` is triggered

initProducer(createDefaultProducerConfig().setBatchingMaxMessages(static_cast<unsigned>(-1)));

std::mutex mtx;

std::vector<Result> results;

auto saveResult = [&mtx, &results](Result result) {

std::lock_guard<std::mutex> lock(mtx);

results.emplace_back(result);

};

auto sendAsync = [saveResult, this](const std::string& key, const std::string& value) {

producer_.sendAsync(MessageBuilder().setOrderingKey(key).setContent(value).build(),

[saveResult](Result result, const MessageId& id) { saveResult(result); });

};

constexpr int numKeys = 10;

for (int i = 0; i < numKeys; i++) {

sendAsync("key-" + std::to_string(i), "value");

}

ASSERT_EQ(ResultOk, producer_.close());

// After close() completed, all callbacks should have failed with ResultAlreadyClosed

std::lock_guard<std::mutex> lock(mtx);

ASSERT_EQ(results.size(), numKeys);

for (int i = 0; i < numKeys; i++) {

ASSERT_EQ(results[i], ResultAlreadyClosed) << " results[" << i << "] is " << results[i];

}