MessageIdTest.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  The ASF licenses this file
 * to you under the Apache License, Version 2.0 (the
 * "License"); you may not use this file except in compliance
 * with the License.  You may obtain a copy of the License at
 *
 *   http://www.apache.org/licenses/LICENSE-2.0
 *
 * 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.
 */
package org.apache.pulsar.client.impl;

import static org.mockito.Matchers.any;
import static org.mockito.Matchers.anyLong;
import static org.mockito.Mockito.doAnswer;
import static org.mockito.Mockito.doReturn;
import static org.mockito.Mockito.doThrow;
import static org.mockito.Mockito.spy;
import static org.mockito.Mockito.times;
import static org.mockito.Mockito.verify;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.fail;

import java.lang.reflect.Field;
import java.util.ArrayList;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.Future;
import java.util.concurrent.TimeUnit;

import org.apache.pulsar.broker.service.BrokerTestBase;
import org.apache.pulsar.client.admin.PulsarAdminException;
import org.apache.pulsar.client.api.ClientConfiguration;
import org.apache.pulsar.client.api.Consumer;
import org.apache.pulsar.client.api.Message;
import org.apache.pulsar.client.api.MessageBuilder;
import org.apache.pulsar.client.api.MessageId;
import org.apache.pulsar.client.api.Producer;
import org.apache.pulsar.client.api.ProducerConfiguration;
import org.apache.pulsar.client.api.PulsarClientException;
import org.apache.pulsar.client.impl.ProducerImpl.OpSendMsg;
import org.apache.pulsar.client.impl.conf.ClientConfigurationData;
import org.apache.pulsar.common.api.ByteBufPair;
import org.apache.pulsar.common.api.Commands;
import org.apache.pulsar.common.api.Commands.ChecksumType;
import org.apache.pulsar.common.api.proto.PulsarApi.MessageMetadata;
import org.apache.pulsar.common.api.proto.PulsarApi.MessageMetadata.Builder;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;

import io.netty.buffer.ByteBuf;

public class MessageIdTest extends BrokerTestBase {
    private static final Logger log = LoggerFactory.getLogger(MessageIdTest.class);

    @BeforeMethod
    @Override
    public void setup() throws Exception {
        baseSetup();
    }

    @AfterMethod
    @Override
    protected void cleanup() throws Exception {
        internalCleanup();
    }

    @Test(timeOut = 10000)
    public void producerSendAsync() throws PulsarClientException {
        // 1. Basic Config
        String key = "producerSendAsync";
        final String topicName = "persistent://prop/cluster/namespace/topic-" + key;
        final String subscriptionName = "my-subscription-" + key;
        final String messagePredicate = "my-message-" + key + "-";
        final int numberOfMessages = 30;

        // 2. Create Producer
        Producer producer = pulsarClient.createProducer(topicName);

        // 3. Create Consumer
        Consumer consumer = pulsarClient.subscribe(topicName, subscriptionName);

        // 4. Publish message and get message id
        Set<MessageId> messageIds = new HashSet();
        List<Future<MessageId>> futures = new ArrayList();
        for (int i = 0; i < numberOfMessages; i++) {
            String message = messagePredicate + i;
            futures.add(producer.sendAsync(message.getBytes()));
        }

        MessageIdImpl previousMessageId = null;
        for (Future<MessageId> f : futures) {
            try {
                MessageIdImpl currentMessageId = (MessageIdImpl) f.get();
                if (previousMessageId != null) {
                    Assert.assertTrue(currentMessageId.compareTo(previousMessageId) > 0,
                            "Message Ids should be in ascending order");
                }
                messageIds.add(currentMessageId);
                previousMessageId = currentMessageId;
            } catch (Exception e) {
                Assert.fail("Failed to publish message, Exception: " + e.getMessage());
            }
        }

        // 4. Check if message Ids are correct
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), numberOfMessages, "Not all messages published successfully");

        for (int i = 0; i < numberOfMessages; i++) {
            Message message = consumer.receive();
            Assert.assertEquals(new String(message.getData()), messagePredicate + i);
            MessageId messageId = message.getMessageId();
            Assert.assertTrue(messageIds.remove(messageId), "Failed to receive message");
        }
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), 0, "Not all messages received successfully");
        consumer.unsubscribe();
    }

    @Test(timeOut = 10000)
    public void producerSend() throws PulsarClientException {
        // 1. Basic Config
        String key = "producerSend";
        final String topicName = "persistent://prop/cluster/namespace/topic-" + key;
        final String subscriptionName = "my-subscription-" + key;
        final String messagePredicate = "my-message-" + key + "-";
        final int numberOfMessages = 30;

        // 2. Create Producer
        Producer producer = pulsarClient.createProducer(topicName);

        // 3. Create Consumer
        Consumer consumer = pulsarClient.subscribe(topicName, subscriptionName);

        // 4. Publish message and get message id
        Set<MessageId> messageIds = new HashSet();
        for (int i = 0; i < numberOfMessages; i++) {
            String message = messagePredicate + i;
            messageIds.add(producer.send(message.getBytes()));
        }

        // 4. Check if message Ids are correct
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), numberOfMessages, "Not all messages published successfully");

        for (int i = 0; i < numberOfMessages; i++) {
            Assert.assertTrue(messageIds.remove(consumer.receive().getMessageId()), "Failed to receive Message");
        }
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), 0, "Not all messages received successfully");
        consumer.unsubscribe();
        ;
    }

    @Test(timeOut = 10000)
    public void partitionedProducerSendAsync() throws PulsarClientException, PulsarAdminException {
        // 1. Basic Config
        String key = "partitionedProducerSendAsync";
        final String topicName = "persistent://prop/cluster/namespace/topic-" + key;
        final String subscriptionName = "my-subscription-" + key;
        final String messagePredicate = "my-message-" + key + "-";
        final int numberOfMessages = 30;
        int numberOfPartitions = 3;
        admin.persistentTopics().createPartitionedTopic(topicName, numberOfPartitions);

        // 2. Create Producer
        Producer producer = pulsarClient.createProducer(topicName);

        // 3. Create Consumer
        Consumer consumer = pulsarClient.subscribe(topicName, subscriptionName);

        // 4. Publish message and get message id
        Set<MessageId> messageIds = new HashSet();
        Set<Future<MessageId>> futures = new HashSet();
        for (int i = 0; i < numberOfMessages; i++) {
            String message = messagePredicate + i;
            futures.add(producer.sendAsync(message.getBytes()));
        }

        futures.forEach(f -> {
            try {
                messageIds.add(f.get());
            } catch (Exception e) {
                Assert.fail("Failed to publish message, Exception: " + e.getMessage());
            }
        });

        // 4. Check if message Ids are correct
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), numberOfMessages, "Not all messages published successfully");

        for (int i = 0; i < numberOfMessages; i++) {
            MessageId messageId = consumer.receive().getMessageId();
            log.info("Message ID Received = " + messageId);
            Assert.assertTrue(messageIds.remove(messageId), "Failed to receive Message");
        }
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), 0, "Not all messages received successfully");
        consumer.unsubscribe();
    }

    @Test(timeOut = 10000)
    public void partitionedProducerSend() throws PulsarClientException, PulsarAdminException {
        // 1. Basic Config
        String key = "partitionedProducerSend";
        final String topicName = "persistent://prop/cluster/namespace/topic-" + key;
        final String subscriptionName = "my-subscription-" + key;
        final String messagePredicate = "my-message-" + key + "-";
        final int numberOfMessages = 30;
        int numberOfPartitions = 7;
        admin.persistentTopics().createPartitionedTopic(topicName, numberOfPartitions);

        // 2. Create Producer
        Producer producer = pulsarClient.createProducer(topicName);

        // 3. Create Consumer
        Consumer consumer = pulsarClient.subscribe(topicName, subscriptionName);

        // 4. Publish message and get message id
        Set<MessageId> messageIds = new HashSet();
        for (int i = 0; i < numberOfMessages; i++) {
            String message = messagePredicate + i;
            messageIds.add(producer.send(message.getBytes()));
        }

        // 4. Check if message Ids are correct
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), numberOfMessages, "Not all messages published successfully");

        for (int i = 0; i < numberOfMessages; i++) {
            Assert.assertTrue(messageIds.remove(consumer.receive().getMessageId()), "Failed to receive Message");
        }
        log.info("Message IDs = " + messageIds);
        Assert.assertEquals(messageIds.size(), 0, "Not all messages received successfully");
        // TODO - this statement causes the broker to hang - need to look into
        // it
        // consumer.unsubscribe();;
    }

    /**
     * Verifies: different versions of broker-deployment (one broker understands Checksum and other
     * doesn't in that case remove checksum before sending to broker-2)
     *
     * client first produce message with checksum and then retries to send message due to connection unavailable. But this time, if
     * broker doesn't understand checksum: then client should remove checksum from the message before sending to broker.
     *
     * 1. stop broker
     * 2. client compute checksum and add into message
     * 3. produce 2 messages and corrupt 1 message
     * 4. start broker with lower version (which doesn't support checksum)
     * 5. client reconnects to broker and due to incompatibility of version: removes checksum from message
     * 6. broker doesn't do checksum validation and persist message
     * 7. client receives ack
     *
     * @throws Exception
     */
    @Test
    public void testChecksumVersionComptability() throws Exception {
        final String topicName = "persistent://prop/use/ns-abc/topic1";

        // 1. producer connect
        ProducerImpl prod = (ProducerImpl) pulsarClient.createProducer(topicName);
        ProducerImpl producer = spy(prod);
        // return higher version compare to broker : so, it forces client-producer to remove checksum from payload
        doReturn(producer.brokerChecksumSupportedVersion() + 1).when(producer).brokerChecksumSupportedVersion();
        doAnswer(invocationOnMock -> prod.getState()).when(producer).getState();
        doAnswer(invocationOnMock -> prod.getClientCnx()).when(producer).getClientCnx();
        doAnswer(invocationOnMock -> prod.cnx()).when(producer).cnx();

        Consumer consumer = pulsarClient.subscribe(topicName, "my-sub");

        // Stop the broker, and publishes messages. Messages are accumulated in the producer queue and they're checksums
        // would have already been computed. If we change the message content at that point, it should result in a
        // checksum validation error
        stopBroker();

        // stop timer to auto-reconnect as let spy-Producer connect to broker manually so, spy-producer object can get
        // mock-value from brokerChecksumSupportedVersion
        ((PulsarClientImpl) pulsarClient).timer().stop();

        ClientCnx mockClientCnx = spy(
                new ClientCnx(new ClientConfigurationData(), ((PulsarClientImpl) pulsarClient).eventLoopGroup()));
        doReturn(producer.brokerChecksumSupportedVersion() - 1).when(mockClientCnx).getRemoteEndpointProtocolVersion();
        prod.setClientCnx(mockClientCnx);

        Message msg1 = MessageBuilder.create().setContent("message-1".getBytes()).build();
        CompletableFuture<MessageId> future1 = producer.sendAsync(msg1);

        Message msg2 = MessageBuilder.create().setContent("message-2".getBytes()).build();
        CompletableFuture<MessageId> future2 = producer.sendAsync(msg2);

        // corrupt the message
        msg2.getData()[msg2.getData().length - 1] = '3'; // new content would be 'message-3'

        prod.setClientCnx(null);

        // Restart the broker to have the messages published
        startBroker();

        // grab broker connection with mocked producer which has higher version compare to broker
        prod.grabCnx();

        try {
            // it should not fail: as due to unsupported version of broker: client removes checksum and broker should
            // ignore the checksum validation
            future1.get();
            future2.get();
        } catch (Exception e) {
            e.printStackTrace();
            fail("Broker shouldn't verify checksum for corrupted message and it shouldn't fail");
        }

        ((ConsumerImpl) consumer).grabCnx();
        // We should only receive msg1
        Message msg = consumer.receive(1, TimeUnit.SECONDS);
        assertEquals(new String(msg.getData()), "message-1");
        msg = consumer.receive(1, TimeUnit.SECONDS);
        assertEquals(new String(msg.getData()), "message-3");

    }

    @Test
    public void testChecksumReconnection() throws Exception {
        final String topicName = "persistent://prop/use/ns-abc/topic1";

        // 1. producer connect
        ProducerImpl prod = (ProducerImpl) pulsarClient.createProducer(topicName);
        ProducerImpl producer = spy(prod);
        // mock: broker-doesn't support checksum (remote_version < brokerChecksumSupportedVersion) so, it forces
        // client-producer to perform checksum-strip from msg at reconnection
        doReturn(producer.brokerChecksumSupportedVersion() + 1).when(producer).brokerChecksumSupportedVersion();
        doAnswer(invocationOnMock -> prod.getState()).when(producer).getState();
        doAnswer(invocationOnMock -> prod.getClientCnx()).when(producer).getClientCnx();
        doAnswer(invocationOnMock -> prod.cnx()).when(producer).cnx();

        Consumer consumer = pulsarClient.subscribe(topicName, "my-sub");

        stopBroker();

        // stop timer to auto-reconnect as let spy-Producer connect to broker
        // manually so, spy-producer object can get
        // mock-value from brokerChecksumSupportedVersion
        ((PulsarClientImpl) pulsarClient).timer().stop();

        // set clientCnx mock to get non-checksum supported version
        ClientCnx mockClientCnx = spy(
                new ClientCnx(new ClientConfigurationData(), ((PulsarClientImpl) pulsarClient).eventLoopGroup()));
        doReturn(producer.brokerChecksumSupportedVersion() - 1).when(mockClientCnx).getRemoteEndpointProtocolVersion();
        prod.setClientCnx(mockClientCnx);

        Message msg1 = MessageBuilder.create().setContent("message-1".getBytes()).build();
        CompletableFuture<MessageId> future1 = producer.sendAsync(msg1);

        Message msg2 = MessageBuilder.create().setContent("message-2".getBytes()).build();
        CompletableFuture<MessageId> future2 = producer.sendAsync(msg2);

        // corrupt the message
        msg2.getData()[msg2.getData().length - 1] = '3'; // new content would be
                                                         // 'message-3'
        // unset mock
        prod.setClientCnx(null);

        // Restart the broker to have the messages published
        startBroker();

        // grab broker connection with mocked producer which has higher version
        // compare to broker
        prod.grabCnx();

        try {
            // it should not fail: as due to unsupported version of broker:
            // client removes checksum and broker should
            // ignore the checksum validation
            future1.get(10, TimeUnit.SECONDS);
            future2.get(10, TimeUnit.SECONDS);
        } catch (Exception e) {
            e.printStackTrace();
            fail("Broker shouldn't verify checksum for corrupted message and it shouldn't fail");
        }

        ((ConsumerImpl) consumer).grabCnx();
        // We should only receive msg1
        Message msg = consumer.receive(1, TimeUnit.SECONDS);
        assertEquals(new String(msg.getData()), "message-1");
        msg = consumer.receive(1, TimeUnit.SECONDS);
        assertEquals(new String(msg.getData()), "message-3");

    }


    /**
     * Verifies: if message is corrupted before sending to broker and if broker gives checksum error: then
     * 1. Client-Producer recomputes checksum with modified data
     * 2. Retry message-send again
     * 3. Broker verifies checksum
     * 4. client receives send-ack success
     *
     * @throws Exception
     */
    @Test
    public void testCorruptMessageRemove() throws Exception {

        final String topicName = "persistent://prop/use/ns-abc/retry-topic";

        ProducerConfiguration config = new ProducerConfiguration();
        config.setSendTimeout(10, TimeUnit.MINUTES);
        // 1. producer connect
        Producer prod = pulsarClient.createProducer(topicName, config);
        ProducerImpl producer = spy((ProducerImpl) prod);
        Field producerIdField = ProducerImpl.class.getDeclaredField("producerId");
        producerIdField.setAccessible(true);
        long producerId = (long) producerIdField.get(producer);
        producer.cnx().registerProducer(producerId, producer); // registered spy ProducerImpl
        Consumer consumer = pulsarClient.subscribe(topicName, "my-sub");

        // 2. Stop the broker, and publishes messages. Messages are accumulated in the producer queue and they're
        // checksums
        // would have already been computed. If we change the message content at that point, it should result in a
        // checksum validation error
        // enable checksum at producer
        stopBroker();

        Message msg = MessageBuilder.create().setContent("message-1".getBytes()).build();
        CompletableFuture<MessageId> future = producer.sendAsync(msg);

        // 3. corrupt the message
        msg.getData()[msg.getData().length - 1] = '2'; // new content would be 'message-3'

        // 4. Restart the broker to have the messages published
        startBroker();

        try {
            future.get();
            fail("send message should have failed with checksum excetion");
        } catch (Exception e) {
            if (e.getCause() instanceof PulsarClientException.ChecksumException) {
                //ok (callback should get checksum exception as message was modified and corrupt)
            } else {
                fail("Callback should have only failed with ChecksumException", e);
            }
        }

        // 5. Verify

        // (5.1) Verify: producer's recoverChecksumError and updateChecksum invoked
        verify(producer, times(1)).recoverChecksumError(any(), anyLong());
        verify(producer, times(1)).verifyLocalBufferIsNotCorrupted(any());


        /**
         * (5.3) verify: ProducerImpl.verifyLocalBufferIsNotCorrupted() => validates if message
         * is corrupt
         */
        MessageImpl msg2 = (MessageImpl) MessageBuilder.create().setContent("message-1".getBytes()).build();
        ByteBuf payload = msg2.getDataBuffer();
        Builder metadataBuilder = ((MessageImpl) msg).getMessageBuilder();
        MessageMetadata msgMetadata = metadataBuilder.setProducerName("test").setSequenceId(1).setPublishTime(10L)
                .build();
        ByteBufPair cmd = Commands.newSend(producerId, 1, 1, ChecksumType.Crc32c, msgMetadata, payload);
        // (a) create OpSendMsg with message-data : "message-1"
        OpSendMsg op = OpSendMsg.create(((MessageImpl) msg), cmd, 1, null);
        // a.verify: as message is not corrupt: no need to update checksum
        assertTrue(producer.verifyLocalBufferIsNotCorrupted(op));
        // (b) corrupt message
        msg2.getData()[msg2.getData().length - 1] = '2'; // new content would be 'message-2'
        // b. verify: as message is corrupt: update checksum
        assertFalse(producer.verifyLocalBufferIsNotCorrupted(op));

        assertEquals(producer.getPendingQueueSize(), 0);

        // [2] test-recoverChecksumError functionality
        stopBroker();
        MessageImpl msg1 = (MessageImpl) MessageBuilder.create().setContent("message-1".getBytes()).build();
        future = producer.sendAsync(msg1);
        ClientCnx cnx = spy(
                new ClientCnx(new ClientConfigurationData(), ((PulsarClientImpl) pulsarClient).eventLoopGroup()));
        String exc = "broker is already stopped";
        // when client-try to recover checksum by resending to broker: throw exception as broker is stopped
        doThrow(new IllegalStateException(exc)).when(cnx).ctx();
        try {
            producer.recoverChecksumError(cnx, 1);
            fail("it should call : resendMessages() => which should throw above mocked exception");
        }catch(IllegalStateException e) {
            assertEquals(exc, e.getMessage());
        }

        producer.close();
        consumer.close();
        producer = null; // clean reference of mocked producer
    }

}