/**
 * 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.api;

import com.google.common.collect.Sets;

import java.util.concurrent.CountDownLatch;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.pulsar.broker.service.Dispatcher;
import org.apache.pulsar.broker.service.persistent.DispatchRateLimiter;
import org.apache.pulsar.broker.service.persistent.PersistentDispatcherMultipleConsumers;
import org.apache.pulsar.broker.service.persistent.PersistentDispatcherSingleActiveConsumer;
import org.apache.pulsar.broker.service.persistent.PersistentSubscription;
import org.apache.pulsar.broker.service.persistent.PersistentTopic;
import org.apache.pulsar.common.policies.data.DispatchRate;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.testng.Assert;
import org.testng.annotations.Test;

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

    /**
     * verify: consumer should not receive all messages due to message-rate throttling
     *
     * @param subscription
     * @throws Exception
     */
    @Test(dataProvider = "subscriptionAndDispatchRateType", timeOut = 5000)
    public void testMessageRateLimitingNotReceiveAllMessages(SubscriptionType subscription,
                                                             DispatchRateType dispatchRateType) throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName = "persistent://" + namespace + "/throttlingBlock";
        final String subName = "my-subscriber-name";

        final int messageRate = 100;
        DispatchRate dispatchRate = null;
        if (DispatchRateType.messageRate.equals(dispatchRateType)) {
            dispatchRate = new DispatchRate(messageRate, -1, 360);
        } else {
            dispatchRate = new DispatchRate(-1, messageRate, 360);
        }

        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);
        // create producer, topic and consumer
        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName).get();

        // enable throttling for nonBacklog consumers
        conf.setDispatchThrottlingOnNonBacklogConsumerEnabled(true);

        final AtomicInteger totalReceived = new AtomicInteger(0);

        Consumer<byte[]> consumer = pulsarClient.newConsumer().topic(topicName).subscriptionName(subName)
            .subscriptionType(subscription).messageListener((c1, msg) -> {
                Assert.assertNotNull(msg, "Message cannot be null");
                String receivedMessage = new String(msg.getData());
                log.debug("Received message [{}] in the listener", receivedMessage);
                totalReceived.incrementAndGet();
            }).subscribe();

        DispatchRateLimiter subRateLimiter = null;
        Dispatcher subDispatcher = topic.getSubscription(subName).getDispatcher();
        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        boolean isMessageRateUpdate = false;
        int retry = 5;
        for (int i = 0; i < retry; i++) {
            if (subRateLimiter.getDispatchRateOnMsg() > 0
                || subRateLimiter.getDispatchRateOnByte() > 0) {
                isMessageRateUpdate = true;
                break;
            } else {
                if (i != retry - 1) {
                    Thread.sleep(100);
                }
            }
        }
        Assert.assertTrue(isMessageRateUpdate);
        Assert.assertEquals(admin.namespaces().getSubscriptionDispatchRate(namespace), dispatchRate);

        int numMessages = 500;
        // Asynchronously produce messages
        for (int i = 0; i < numMessages; i++) {
            producer.send(new byte[80]);
        }

        // consumer should not have received all published message due to message-rate throttling
        Assert.assertTrue(totalReceived.get() < messageRate * 2);

        consumer.close();
        producer.close();
        log.info("-- Exiting {} test --", methodName);
    }

    /**
     * verify rate-limiting should throttle message-dispatching based on message-rate
     *
     * <pre>
     *  1. dispatch-msg-rate = 10 msg/sec
     *  2. send 30 msgs
     *  3. it should take up to 2 second to receive all messages
     * </pre>
     *
     * @param subscription
     * @throws Exception
     */
    @Test(dataProvider = "subscriptions", timeOut = 5000)
    public void testMessageRateLimitingReceiveAllMessagesAfterThrottling(SubscriptionType subscription)
        throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName = "persistent://" + namespace + "/throttlingAll";
        final String subName = "my-subscriber-name";

        final int messageRate = 10;
        DispatchRate dispatchRate = new DispatchRate(messageRate, -1, 1);
        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);
        final int numProducedMessages = 30;
        final CountDownLatch latch = new CountDownLatch(numProducedMessages);
        final AtomicInteger totalReceived = new AtomicInteger(0);
        // enable throttling for nonBacklog consumers
        conf.setDispatchThrottlingOnNonBacklogConsumerEnabled(true);

        Consumer<byte[]> consumer = pulsarClient.newConsumer().topic(topicName).subscriptionName(subName)
            .subscriptionType(subscription).messageListener((c1, msg) -> {
                Assert.assertNotNull(msg, "Message cannot be null");
                String receivedMessage = new String(msg.getData());
                log.debug("Received message [{}] in the listener", receivedMessage);
                totalReceived.incrementAndGet();
                latch.countDown();
            }).subscribe();

        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName).get();

        DispatchRateLimiter subRateLimiter = null;
        Dispatcher subDispatcher = topic.getSubscription(subName).getDispatcher();
        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        boolean isMessageRateUpdate = false;
        int retry = 5;
        for (int i = 0; i < retry; i++) {
            if (subRateLimiter.getDispatchRateOnMsg() > 0
                || subRateLimiter.getDispatchRateOnByte() > 0) {
                isMessageRateUpdate = true;
                break;
            } else {
                if (i != retry - 1) {
                    Thread.sleep(100);
                }
            }
        }
        Assert.assertTrue(isMessageRateUpdate);
        Assert.assertEquals(admin.namespaces().getSubscriptionDispatchRate(namespace), dispatchRate);

        long start = System.currentTimeMillis();
        // Asynchronously produce messages
        for (int i = 0; i < numProducedMessages; i++) {
            final String message = "my-message-" + i;
            producer.send(message.getBytes());
        }
        latch.await();
        Assert.assertEquals(totalReceived.get(), numProducedMessages);
        long end = System.currentTimeMillis();

        Assert.assertTrue((end - start) >= 2000);

        consumer.close();
        producer.close();
        log.info("-- Exiting {} test --", methodName);
    }

    /**
     * verify rate-limiting should throttle message-dispatching based on byte-rate
     *
     * <pre>
     *  1. dispatch-byte-rate = 100 bytes/sec
     *  2. send 30 msgs : each with 10 byte
     *  3. it should take up to 2 second to receive all messages
     * </pre>
     *
     * @param subscription
     * @throws Exception
     */
    @Test(dataProvider = "subscriptions", timeOut = 5000)
    public void testBytesRateLimitingReceiveAllMessagesAfterThrottling(SubscriptionType subscription) throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName = "persistent://" + namespace + "/throttlingAll-" + System.nanoTime();
        final String subName = "my-subscriber-name-" + subscription;

        final int byteRate = 100;
        DispatchRate dispatchRate = new DispatchRate(-1, byteRate, 1);
        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);
        final int numProducedMessages = 30;
        final CountDownLatch latch = new CountDownLatch(numProducedMessages);
        final AtomicInteger totalReceived = new AtomicInteger(0);
        // enable throttling for nonBacklog consumers
        conf.setDispatchThrottlingOnNonBacklogConsumerEnabled(true);

        Consumer<byte[]> consumer = pulsarClient.newConsumer().topic(topicName).subscriptionName(subName)
            .receiverQueueSize(10)
            .subscriptionType(subscription).messageListener((c1, msg) -> {
                Assert.assertNotNull(msg, "Message cannot be null");
                String receivedMessage = new String(msg.getData());
                log.debug("Received message [{}] in the listener", receivedMessage);
                totalReceived.incrementAndGet();
                latch.countDown();
            }).subscribe();

        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName).get();

        DispatchRateLimiter subRateLimiter = null;
        Dispatcher subDispatcher = topic.getSubscription(subName).getDispatcher();
        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        boolean isMessageRateUpdate = false;
        int retry = 5;
        for (int i = 0; i < retry; i++) {
            if (subRateLimiter.getDispatchRateOnMsg() > 0
                || subRateLimiter.getDispatchRateOnByte() > 0) {
                isMessageRateUpdate = true;
                break;
            } else {
                if (i != retry - 1) {
                    Thread.sleep(100);
                }
            }
        }
        Assert.assertTrue(isMessageRateUpdate);
        Assert.assertEquals(admin.namespaces().getSubscriptionDispatchRate(namespace), dispatchRate);

        long start = System.currentTimeMillis();
        // Asynchronously produce messages
        for (int i = 0; i < numProducedMessages; i++) {
            producer.send(new byte[byteRate / 10]);
        }
        latch.await();
        Assert.assertEquals(totalReceived.get(), numProducedMessages, 10);
        long end = System.currentTimeMillis();
        log.info("-- end - start: {} ", end - start);

        // first 10 messages, which equals receiverQueueSize, will not wait.
        Assert.assertTrue((end - start) >= 2000);

        consumer.close();
        producer.close();
        log.info("-- Exiting {} test --", methodName);
    }

    /**
     * verify message-rate on multiple consumers with shared-subscription
     *
     * @throws Exception
     */
    @Test(timeOut = 5000)
    public void testRateLimitingMultipleConsumers() throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName = "persistent://" + namespace + "/throttlingMultipleConsumers";
        final String subName = "my-subscriber-name";

        final int messageRate = 5;
        DispatchRate dispatchRate = new DispatchRate(messageRate, -1, 360);
        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);

        final int numProducedMessages = 500;
        final AtomicInteger totalReceived = new AtomicInteger(0);
        // enable throttling for nonBacklog consumers
        conf.setDispatchThrottlingOnNonBacklogConsumerEnabled(true);

        ConsumerBuilder<byte[]> consumerBuilder = pulsarClient.newConsumer().topic(topicName)
            .subscriptionName(subName).subscriptionType(SubscriptionType.Shared).messageListener((c1, msg) -> {
                Assert.assertNotNull(msg, "Message cannot be null");
                String receivedMessage = new String(msg.getData());
                log.debug("Received message [{}] in the listener", receivedMessage);
                totalReceived.incrementAndGet();
            });
        Consumer<byte[]> consumer1 = consumerBuilder.subscribe();
        Consumer<byte[]> consumer2 = consumerBuilder.subscribe();
        Consumer<byte[]> consumer3 = consumerBuilder.subscribe();
        Consumer<byte[]> consumer4 = consumerBuilder.subscribe();
        Consumer<byte[]> consumer5 = consumerBuilder.subscribe();

        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName).get();

        DispatchRateLimiter subRateLimiter = null;
        Dispatcher subDispatcher = topic.getSubscription(subName).getDispatcher();
        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        boolean isMessageRateUpdate = false;
        int retry = 5;
        for (int i = 0; i < retry; i++) {
            if (subRateLimiter.getDispatchRateOnMsg() > 0
                || subRateLimiter.getDispatchRateOnByte() > 0) {
                isMessageRateUpdate = true;
                break;
            } else {
                if (i != retry - 1) {
                    Thread.sleep(100);
                }
            }
        }
        Assert.assertTrue(isMessageRateUpdate);
        Assert.assertEquals(admin.namespaces().getSubscriptionDispatchRate(namespace), dispatchRate);

        // Asynchronously produce messages
        for (int i = 0; i < numProducedMessages; i++) {
            final String message = "my-message-" + i;
            producer.send(message.getBytes());
        }

        // it can make sure that consumer had enough time to consume message but couldn't consume due to throttling
        Thread.sleep(500);

        // consumer should not have received all published message due to message-rate throttling
        Assert.assertNotEquals(totalReceived.get(), numProducedMessages);

        consumer1.close();
        consumer2.close();
        consumer3.close();
        consumer4.close();
        consumer5.close();
        producer.close();
        log.info("-- Exiting {} test --", methodName);
    }


    @Test(dataProvider = "subscriptions", timeOut = 5000)
    public void testClusterRateLimitingConfiguration(SubscriptionType subscription) throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName = "persistent://" + namespace + "/throttlingBlock";
        final String subName = "my-subscriber-name";

        final int messageRate = 5;
        // enable throttling for nonBacklog consumers
        conf.setDispatchThrottlingOnNonBacklogConsumerEnabled(true);

        int initValue = pulsar.getConfiguration().getDispatchThrottlingRatePerSubscriptionInMsg();
        // (1) Update message-dispatch-rate limit
        admin.brokers().updateDynamicConfiguration("dispatchThrottlingRatePerSubscriptionInMsg",
            Integer.toString(messageRate));
        // sleep incrementally as zk-watch notification is async and may take some time
        for (int i = 0; i < 5; i++) {
            if (pulsar.getConfiguration().getDispatchThrottlingRatePerSubscriptionInMsg() == initValue) {
                Thread.sleep(50 + (i * 10));
            }
        }
        Assert.assertNotEquals(pulsar.getConfiguration().getDispatchThrottlingRatePerSubscriptionInMsg(), initValue);

        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));
        // create producer and topic
        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName).get();
        int numMessages = 500;

        final AtomicInteger totalReceived = new AtomicInteger(0);

        Consumer<byte[]> consumer = pulsarClient.newConsumer().topic(topicName).subscriptionName(subName)
            .subscriptionType(subscription).messageListener((c1, msg) -> {
                Assert.assertNotNull(msg, "Message cannot be null");
                String receivedMessage = new String(msg.getData());
                log.debug("Received message [{}] in the listener", receivedMessage);
                totalReceived.incrementAndGet();
            }).subscribe();

        // Asynchronously produce messages
        for (int i = 0; i < numMessages; i++) {
            final String message = "my-message-" + i;
            producer.send(message.getBytes());
        }

        // it can make sure that consumer had enough time to consume message but couldn't consume due to throttling
        Thread.sleep(500);

        // consumer should not have received all published message due to message-rate throttling
        Assert.assertNotEquals(totalReceived.get(), numMessages);

        consumer.close();
        producer.close();
        pulsar.getConfiguration().setDispatchThrottlingRatePerSubscriptionInMsg(initValue);
        log.info("-- Exiting {} test --", methodName);
    }

    /**
     * <pre>
     * It verifies that cluster-throttling value gets considered when namespace-policy throttling is disabled.
     *
     *  1. Update cluster-throttling-config: topic rate-limiter has cluster-config
     *  2. Update namespace-throttling-config: topic rate-limiter has namespace-config
     *  3. Disable namespace-throttling-config: topic rate-limiter has cluster-config
     *  4. Create new topic with disable namespace-config and enabled cluster-config: it takes cluster-config
     *
     * </pre>
     *
     * @throws Exception
     */
    @Test
    public void testClusterPolicyOverrideConfiguration() throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName1 = "persistent://" + namespace + "/throttlingOverride1";
        final String topicName2 = "persistent://" + namespace + "/throttlingOverride2";
        final String subName1 = "my-subscriber-name1";
        final String subName2 = "my-subscriber-name2";

        final int clusterMessageRate = 100;
        // enable throttling for nonBacklog consumers
        conf.setDispatchThrottlingOnNonBacklogConsumerEnabled(true);

        int initValue = pulsar.getConfiguration().getDispatchThrottlingRatePerSubscriptionInMsg();
        // (1) Update message-dispatch-rate limit
        admin.brokers().updateDynamicConfiguration("dispatchThrottlingRatePerSubscriptionInMsg",
            Integer.toString(clusterMessageRate));
        // sleep incrementally as zk-watch notification is async and may take some time
        for (int i = 0; i < 5; i++) {
            if (pulsar.getConfiguration().getDispatchThrottlingRatePerSubscriptionInMsg() == initValue) {
                Thread.sleep(50 + (i * 10));
            }
        }
        Assert.assertNotEquals(pulsar.getConfiguration().getDispatchThrottlingRatePerSubscriptionInMsg(), initValue);

        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));

        // create producer and topic
        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName1).create();
        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName1).get();

        Consumer<byte[]> consumer1 = pulsarClient.newConsumer().topic(topicName1).subscriptionName(subName1)
            .subscribe();

        DispatchRateLimiter subRateLimiter = null;
        Dispatcher subDispatcher = topic.getSubscription(subName1).getDispatcher();
        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        // (1) Update dispatch rate on cluster-config update
        Assert.assertEquals(clusterMessageRate, subRateLimiter.getDispatchRateOnMsg());

        // (2) Update namespace throttling limit
        int nsMessageRate = 500;
        DispatchRate dispatchRate = new DispatchRate(nsMessageRate, 0, 1);
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);

        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        for (int i = 0; i < 5; i++) {
            if (subRateLimiter.getDispatchRateOnMsg() != nsMessageRate) {
                Thread.sleep(50 + (i * 10));
            }
        }
        Assert.assertEquals(nsMessageRate, subRateLimiter.getDispatchRateOnMsg());

        // (3) Disable namespace throttling limit will force to take cluster-config
        dispatchRate = new DispatchRate(0, 0, 1);
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);
        for (int i = 0; i < 5; i++) {
            if (subRateLimiter.getDispatchRateOnMsg() == nsMessageRate) {
                Thread.sleep(50 + (i * 10));
            }
        }
        Assert.assertEquals(clusterMessageRate, subRateLimiter.getDispatchRateOnMsg());

        // (5) Namespace throttling is disabled so, new topic should take cluster throttling limit
        Producer<byte[]> producer2 = pulsarClient.newProducer().topic(topicName2).create();
        PersistentTopic topic2 = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName2).get();
        Consumer<byte[]> consumer2 = pulsarClient.newConsumer().topic(topicName2).subscriptionName(subName2)
            .subscribe();

        subDispatcher = topic2.getSubscription(subName2).getDispatcher();
        if (subDispatcher instanceof PersistentDispatcherMultipleConsumers) {
            subRateLimiter = ((PersistentDispatcherMultipleConsumers) subDispatcher).getRateLimiter().get();
        } else if (subDispatcher instanceof PersistentDispatcherSingleActiveConsumer) {
            subRateLimiter = ((PersistentDispatcherSingleActiveConsumer) subDispatcher).getRateLimiter().get();
        } else {
            Assert.fail("Should only have PersistentDispatcher in this test");
        }

        Assert.assertEquals(clusterMessageRate, subRateLimiter.getDispatchRateOnMsg());

        producer.close();
        producer2.close();

        log.info("-- Exiting {} test --", methodName);
    }

    @Test(dataProvider = "subscriptions", timeOut = 10000)
    public void testClosingRateLimiter(SubscriptionType subscription) throws Exception {
        log.info("-- Starting {} test --", methodName);

        final String namespace = "my-property/throttling_ns";
        final String topicName = "persistent://" + namespace + "/closingSubRateLimiter" + subscription.name();
        final String subName = "mySubscription" + subscription.name();

        DispatchRate dispatchRate = new DispatchRate(10, 1024, 1);
        admin.namespaces().createNamespace(namespace, Sets.newHashSet("test"));
        admin.namespaces().setSubscriptionDispatchRate(namespace, dispatchRate);

        Producer<byte[]> producer = pulsarClient.newProducer().topic(topicName).create();
        Consumer<byte[]> consumer = pulsarClient.newConsumer().topic(topicName).subscriptionName(subName)
                .subscriptionType(subscription).subscribe();

        PersistentTopic topic = (PersistentTopic) pulsar.getBrokerService().getOrCreateTopic(topicName).get();
        PersistentSubscription sub = topic.getSubscription(subName);

        final int numProducedMessages = 10;

        for (int i = 0; i < numProducedMessages; i++) {
            final String message = "my-message-" + i;
            producer.send(message.getBytes());
        }

        for (int i = 0; i < numProducedMessages; i++) {
            Message<byte[]> msg = consumer.receive();
            consumer.acknowledge(msg);
        }

        Dispatcher dispatcher = sub.getDispatcher();
        Assert.assertTrue(dispatcher.getRateLimiter().isPresent());
        DispatchRateLimiter dispatchRateLimiter = dispatcher.getRateLimiter().get();

        producer.close();
        consumer.close();
        sub.disconnect().get();

        // Make sure that the rate limiter is closed
        Assert.assertEquals(dispatchRateLimiter.getDispatchRateOnMsg(), -1);
        Assert.assertEquals(dispatchRateLimiter.getDispatchRateOnByte(), -1);

        log.info("-- Exiting {} test --", methodName);
    }
}
