[fix] Consumer batch receive will cause data loss. (#228)

* [fix] Consumer batch receive will cause data loss.

* Support popIf on queue.

* Fix code review.

Author: shibd

lib/BlockingQueue.h

@@ -120,6 +120,37 @@ class BlockingQueue {
         return true;
     }
 
+    /**
+     * First peek data to the condition judgment, if true then pop it.
+     *
+     * @param value A reference to the value assigned after pop
+     * @param condition A function that returns true if the value should be popped
+     * @return true if the value was popped, false otherwise.
+     */
+    bool popIf(T& value, std::function<bool(const T& peekValue)> condition) {
+        Lock lock(mutex_);
+
+        if (isEmptyNoMutex() || isClosedNoMutex()) {
+            return false;
+        }
+
+        bool wasFull = isFullNoMutex();
+
+        auto peekValue = queue_.front();
+        if (condition(peekValue)) {
+            value = peekValue;
+            queue_.pop_front();
+            lock.unlock();
+            if (wasFull) {
+                // Notify that an element is popped
+                queueFullCondition.notify_all();
+            }
+            return true;
+        } else {
+            return false;
+        }
+    }
+
     // Check the 1st element of the queue
     bool peek(T& value) {
         Lock lock(mutex_);

lib/ConsumerImpl.cc

@@ -665,10 +665,11 @@ void ConsumerImpl::executeNotifyCallback(Message& msg) {
 void ConsumerImpl::notifyBatchPendingReceivedCallback(const BatchReceiveCallback& callback) {
     auto messages = std::make_shared<MessagesImpl>(batchReceivePolicy_.getMaxNumMessages(),
                                                    batchReceivePolicy_.getMaxNumBytes());
-    Message peekMsg;
-    while (incomingMessages_.pop(peekMsg, std::chrono::milliseconds(0)) && messages->canAdd(peekMsg)) {
-        messageProcessed(peekMsg);
-        Message interceptMsg = interceptors_->beforeConsume(Consumer(shared_from_this()), peekMsg);
+    Message msg;
+    while (incomingMessages_.popIf(
+        msg, [&messages](const Message& peekMsg) { return messages->canAdd(peekMsg); })) {
+        messageProcessed(msg);
+        Message interceptMsg = interceptors_->beforeConsume(Consumer(shared_from_this()), msg);
         messages->add(interceptMsg);
     }
     auto self = get_shared_this_ptr();

lib/MultiTopicsConsumerImpl.cc

@@ -1054,10 +1054,11 @@ bool MultiTopicsConsumerImpl::hasEnoughMessagesForBatchReceive() const {
 void MultiTopicsConsumerImpl::notifyBatchPendingReceivedCallback(const BatchReceiveCallback& callback) {
     auto messages = std::make_shared<MessagesImpl>(batchReceivePolicy_.getMaxNumMessages(),
                                                    batchReceivePolicy_.getMaxNumBytes());
-    Message peekMsg;
-    while (incomingMessages_.pop(peekMsg, std::chrono::milliseconds(0)) && messages->canAdd(peekMsg)) {
-        messageProcessed(peekMsg);
-        messages->add(peekMsg);
+    Message msg;
+    while (incomingMessages_.popIf(
+        msg, [&messages](const Message& peekMsg) { return messages->canAdd(peekMsg); })) {
+        messageProcessed(msg);
+        messages->add(msg);
     }
     auto weakSelf = weak_from_this();
     listenerExecutor_->postWork([weakSelf, callback, messages]() {

lib/UnboundedBlockingQueue.h

@@ -87,6 +87,30 @@ class UnboundedBlockingQueue {
         return true;
     }
 
+    /**
+     * First peek data to the condition judgment, if true then pop it.
+     *
+     * @param value A reference to the value assigned after pop
+     * @param condition A function that returns true if the value should be popped
+     * @return true if the value was popped, false otherwise.
+     */
+    bool popIf(T& value, std::function<bool(const T& peekValue)> condition) {
+        Lock lock(mutex_);
+
+        if (isEmptyNoMutex() || isClosedNoMutex()) {
+            return false;
+        }
+
+        auto peekValue = queue_.front();
+        if (condition(peekValue)) {
+            value = peekValue;
+            queue_.pop_front();
+            return true;
+        } else {
+            return false;
+        }
+    }
+
     // Check the 1st element of the queue
     bool peek(T& value) {
         Lock lock(mutex_);

tests/BasicEndToEndTest.cc

@@ -4084,12 +4084,11 @@ void testBatchReceive(bool multiConsumer) {
     Client client(lookupUrl);
 
     std::string uniqueChunk = unique_str();
-    std::string topicName = "persistent://public/default/test-batch-receive" + uniqueChunk;
+    std::string topicName = "test-batch-receive" + uniqueChunk + std::to_string(multiConsumer);
 
     if (multiConsumer) {
         // call admin api to make it partitioned
-        std::string url =
-            adminUrl + "admin/v2/persistent/public/default/test-batch-receive" + uniqueChunk + "/partitions";
+        std::string url = adminUrl + "admin/v2/persistent/public/default/" + topicName + "/partitions";
         int res = makePutRequest(url, "5");
         LOG_INFO("res = " << res);
         ASSERT_FALSE(res != 204 && res != 409);
@@ -4107,8 +4106,9 @@ void testBatchReceive(bool multiConsumer) {
     Consumer consumer;
     ConsumerConfiguration consumerConfig;
     // when receiver queue size > maxNumMessages, use receiver queue size.
+    const int batchReceiveMaxNumMessages = 10;
     consumerConfig.setBatchReceivePolicy(BatchReceivePolicy(1000, -1, -1));
-    consumerConfig.setReceiverQueueSize(10);
+    consumerConfig.setReceiverQueueSize(batchReceiveMaxNumMessages);
     consumerConfig.setProperty("consumer-name", "test-consumer-name");
     consumerConfig.setProperty("consumer-id", "test-consumer-id");
     Promise<Result, Consumer> consumerPromise;
@@ -4120,32 +4120,41 @@ void testBatchReceive(bool multiConsumer) {
 
     // sync batch receive test
     std::string prefix = "batch-receive-msg";
-    int numOfMessages = 10;
+    int numOfMessages = 100;
     for (int i = 0; i < numOfMessages; i++) {
         std::string messageContent = prefix + std::to_string(i);
         Message msg = MessageBuilder().setContent(messageContent).build();
-        producer.send(msg);
+        producer.sendAsync(msg, NULL);
+    }
+    ASSERT_EQ(ResultOk, producer.flush());
+    for (int i = 0; i < numOfMessages / batchReceiveMaxNumMessages; i++) {
+        Messages messages;
+        Result receive = consumer.batchReceive(messages);
+        ASSERT_EQ(receive, ResultOk);
+        ASSERT_EQ(messages.size(), batchReceiveMaxNumMessages);
+        for (const auto &item : messages) {
+            consumer.acknowledge(item);
+        }
     }
 
-    Messages messages;
-    Result receive = consumer.batchReceive(messages);
-    ASSERT_EQ(receive, ResultOk);
-    ASSERT_EQ(messages.size(), numOfMessages);
-
     // async batch receive test
-    Latch latch(1);
-    BatchReceiveCallback batchReceiveCallback = [&latch, numOfMessages](Result result, Messages messages) {
-        ASSERT_EQ(result, ResultOk);
-        ASSERT_EQ(messages.size(), numOfMessages);
-        latch.countdown();
-    };
-    consumer.batchReceiveAsync(batchReceiveCallback);
     for (int i = 0; i < numOfMessages; i++) {
         std::string messageContent = prefix + std::to_string(i);
         Message msg = MessageBuilder().setContent(messageContent).build();
-        producer.send(msg);
+        producer.sendAsync(msg, NULL);
+    }
+    ASSERT_EQ(ResultOk, producer.flush());
+    for (int i = 0; i < numOfMessages / batchReceiveMaxNumMessages; i++) {
+        Latch latch(1);
+        BatchReceiveCallback batchReceiveCallback = [&latch, batchReceiveMaxNumMessages](Result result,
+                                                                                         Messages messages) {
+            ASSERT_EQ(result, ResultOk);
+            ASSERT_EQ(messages.size(), batchReceiveMaxNumMessages);
+            latch.countdown();
+        };
+        consumer.batchReceiveAsync(batchReceiveCallback);
+        ASSERT_TRUE(latch.wait(std::chrono::seconds(1)));
     }
-    ASSERT_TRUE(latch.wait(std::chrono::seconds(10)));
 
     producer.close();
     consumer.close();

tests/BlockingQueueTest.cc

@@ -126,6 +126,27 @@ TEST(BlockingQueueTest, testBlockingProducer) {
     ASSERT_EQ(three, queue.size());
 }
 
+TEST(BlockingQueueTest, testPopIf) {
+    size_t size = 5;
+    BlockingQueue<int> queue(size);
+
+    for (int i = 1; i <= size; ++i) {
+        queue.push(i);
+    }
+
+    // Use producer worker to assert popIf will notify queueFull thread.
+    ProducerWorker producerWorker(queue);
+    producerWorker.produce(3);
+
+    int value;
+    for (int i = 1; i <= size; ++i) {
+        ASSERT_TRUE(queue.popIf(value, [&i](const int& peekValue) { return peekValue == i; }));
+    }
+
+    producerWorker.join();
+    ASSERT_EQ(3, queue.size());
+}
+
 TEST(BlockingQueueTest, testBlockingConsumer) {
     size_t size = 5;
     BlockingQueue<int> queue(size);

tests/UnboundedBlockingQueueTest.cc

@@ -112,6 +112,22 @@ TEST(UnboundedBlockingQueueTest, testQueueOperations) {
     ASSERT_FALSE(queue.peek(poppedElement));
 }
 
+TEST(UnboundedBlockingQueueTest, testPopIf) {
+    size_t size = 5;
+    UnboundedBlockingQueue<int> queue(size);
+
+    for (int i = 1; i <= size * 2; ++i) {
+        queue.push(i);
+    }
+
+    int value;
+    for (int i = 1; i <= size; ++i) {
+        ASSERT_TRUE(queue.popIf(value, [&i](const int& peekValue) { return peekValue == i; }));
+    }
+
+    ASSERT_EQ(size, queue.size());
+}
+
 TEST(UnboundedBlockingQueueTest, testBlockingProducer) {
     size_t size = 5;
     UnboundedBlockingQueue<int> queue(size);