AMQ6133PersistJMSRedeliveryTest.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.activemq.bugs;

import static org.junit.Assert.assertEquals;
import static org.junit.Assert.assertNotNull;
import static org.junit.Assert.assertTrue;

import java.io.File;
import java.io.IOException;
import java.util.ArrayList;
import java.util.concurrent.TimeUnit;

import jakarta.jms.BytesMessage;
import jakarta.jms.Connection;
import jakarta.jms.DeliveryMode;
import jakarta.jms.Destination;
import jakarta.jms.JMSException;
import jakarta.jms.Message;
import jakarta.jms.MessageConsumer;
import jakarta.jms.MessageProducer;
import jakarta.jms.Queue;
import jakarta.jms.Session;
import javax.management.MalformedObjectNameException;
import javax.management.ObjectName;

import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.broker.BrokerService;
import org.apache.activemq.broker.jmx.QueueViewMBean;
import org.apache.activemq.broker.region.policy.PolicyEntry;
import org.apache.activemq.broker.region.policy.PolicyMap;
import org.apache.activemq.store.kahadb.KahaDBPersistenceAdapter;
import org.apache.activemq.store.kahadb.MessageDatabase;
import org.apache.activemq.util.Wait;
import org.apache.commons.io.FileUtils;
import org.apache.commons.io.filefilter.TrueFileFilter;
import org.apache.commons.io.filefilter.WildcardFileFilter;
import org.apache.logging.log4j.Level;
import org.apache.logging.log4j.LogManager;
import org.junit.After;
import org.junit.Before;
import org.junit.Test;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.junit.experimental.categories.Category;
import org.apache.activemq.test.annotations.ParallelTest;

/**
 * Test loss of message on index rebuild when presistJMSRedelivered is on.
 */
@Category(ParallelTest.class)
public class AMQ6133PersistJMSRedeliveryTest {

    private static final Logger LOG = LoggerFactory.getLogger(AMQ6133PersistJMSRedeliveryTest.class);
    private static final String QUEUE_NAME = "test.queue";

    private BrokerService broker;

    @Test
    public void testPersistJMSRedeliveredMessageLossOnIndexRebuild() throws Exception {
        sendMessages();
        LOG.info("#### Finished sending messages, test starting. ####");

        long msgCount = getProxyToQueue(QUEUE_NAME).getQueueSize();

        final int ITERATIONS = 3;

        // Force some updates
        for (int i = 0; i < ITERATIONS; ++i) {
            LOG.info("Consumer and Rollback iteration: {}", i);
            consumerAndRollback(i);
        }

        // Allow GC to run at least once.
        TimeUnit.SECONDS.sleep(20);

        restart();

        assertTrue("Queue size should match expected count: " + msgCount, Wait.waitFor(() -> {
            try {
                return getProxyToQueue(QUEUE_NAME).getQueueSize() == msgCount;
            } catch (Exception e) {
               return false;
            }
        }));

        restartWithRecovery(getPersistentDir());

        assertTrue("Queue size should match expected count: " + msgCount, Wait.waitFor(() -> {
            try {
                return getProxyToQueue(QUEUE_NAME).getQueueSize() == msgCount;
            } catch (Exception e) {
                return false;
            }
        }));
    }

    @Before
    public void setup() throws Exception {

        // Investigate loss of messages on message update in store.
        ((org.apache.logging.log4j.core.Logger)LogManager.getLogger(MessageDatabase.class)).setLevel(Level.TRACE);
        createBroker(true);
    }

    @After
    public void tearDown() throws Exception {
        broker.stop();
    }

    private void restart() throws Exception {
        broker.stop();
        broker.waitUntilStopped();

        createBroker(false);
    }

    private void restartWithRecovery(File persistenceDir) throws Exception {
        broker.stop();
        broker.waitUntilStopped();

        // delete the index so that it needs to be rebuilt from replay
        for (File index : FileUtils.listFiles(persistenceDir, new WildcardFileFilter("db.*"), TrueFileFilter.INSTANCE)) {
            FileUtils.deleteQuietly(index);
        }

        createBroker(false);
    }

    private void sendMessages() throws Exception {
        Connection connection = createConnection();
        Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        Destination queue = session.createQueue(QUEUE_NAME);
        Destination retainQueue = session.createQueue(QUEUE_NAME + "-retain");
        MessageProducer producer = session.createProducer(null);

        final byte[] payload = new byte[1000];
        producer.setDeliveryMode(DeliveryMode.PERSISTENT);
        BytesMessage message = session.createBytesMessage();
        message.writeBytes(payload);

        // Build up a set of messages that will be redelivered and updated later.
        while (getLogFileCount() < 3) {
            producer.send(queue, message);
        }

        // Now create some space for files that are retained during the test.
        while (getLogFileCount() < 6) {
            producer.send(retainQueue, message);
        }

        connection.close();
    }

    private void consumerAndRollback(int iteration) throws Exception {
        Connection connection = createConnection();
        Session session = connection.createSession(false, Session.CLIENT_ACKNOWLEDGE);
        Queue queue = session.createQueue(QUEUE_NAME);
        MessageConsumer consumer = session.createConsumer(queue);

        long msgCount = getProxyToQueue(queue.getQueueName()).getQueueSize();

        for (int i = 0; i < msgCount; ++i) {
            Message message = consumer.receive(50000);
            assertNotNull(message);
            if (iteration > 0) {
                assertTrue(message.getJMSRedelivered());
            }
        }

        connection.close();
    }

    private Connection createConnection() throws Exception {
        ActiveMQConnectionFactory factory = new ActiveMQConnectionFactory("vm://localhost?jms.prefetchPolicy.all=0");
        factory.setAlwaysSyncSend(true);
        Connection connection = factory.createConnection();
        connection.start();

        return connection;
    }

    private void createBroker(boolean deleteAllMessages) throws Exception {
        PolicyEntry entry = new PolicyEntry();
        entry.setPersistJMSRedelivered(true);

        PolicyMap policyMap = new PolicyMap();
        policyMap.setDefaultEntry(entry);

        broker = new BrokerService();
        broker.setDeleteAllMessagesOnStartup(deleteAllMessages);
        broker.setPersistent(true);
        broker.setDestinationPolicy(policyMap);

        KahaDBPersistenceAdapter kahaDB = new KahaDBPersistenceAdapter();
        kahaDB.setJournalMaxFileLength(128 * 1024);
        kahaDB.setCleanupInterval(8*1000);

        broker.setPersistenceAdapter(kahaDB);
        broker.getSystemUsage().getStoreUsage().setLimit(7*1024*1024);
        broker.start();
        broker.waitUntilStarted();
    }

    private int getLogFileCount() throws Exception {
        return new ArrayList<File>(
                FileUtils.listFiles(getPersistentDir(),
                    new WildcardFileFilter("*.log"), TrueFileFilter.INSTANCE)).size();
    }

    private File getPersistentDir() throws IOException {
        return broker.getPersistenceAdapter().getDirectory();
    }

    protected QueueViewMBean getProxyToQueue(String name) throws MalformedObjectNameException, JMSException {
        ObjectName queueViewMBeanName = new ObjectName("org.apache.activemq:type=Broker,brokerName=localhost,destinationType=Queue,destinationName="+name);
        QueueViewMBean proxy = (QueueViewMBean) broker.getManagementContext()
                .newProxyInstance(queueViewMBeanName, QueueViewMBean.class, true);
        return proxy;
    }
}