content: /*
   * 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.artemis.tests.integration.cluster.crossprotocol;
  
  import static org.junit.jupiter.api.Assertions.assertNotNull;
  import static org.junit.jupiter.api.Assertions.assertNull;
  import static org.junit.jupiter.api.Assertions.fail;
  
  import java.lang.invoke.MethodHandles;
  import java.util.Arrays;
  import java.util.Collection;
  
  import javax.jms.Connection;
  import javax.jms.ConnectionFactory;
  import javax.jms.JMSException;
  import javax.jms.Message;
  import javax.jms.MessageConsumer;
  import javax.jms.MessageProducer;
  import javax.jms.Session;
  
  import org.apache.activemq.artemis.api.core.ActiveMQException;
  import org.apache.activemq.artemis.api.core.QueueConfiguration;
  import org.apache.activemq.artemis.api.core.RoutingType;
  import org.apache.activemq.artemis.api.core.SimpleString;
  import org.apache.activemq.artemis.api.core.client.ActiveMQClient;
  import org.apache.activemq.artemis.api.core.client.ClientMessage;
  import org.apache.activemq.artemis.api.core.client.ClientProducer;
  import org.apache.activemq.artemis.api.core.client.ClientSession;
  import org.apache.activemq.artemis.api.core.client.ClientSessionFactory;
  import org.apache.activemq.artemis.api.core.client.ServerLocator;
  import org.apache.activemq.artemis.core.config.impl.ConfigurationImpl;
  import org.apache.activemq.artemis.core.protocol.openwire.OpenWireProtocolManagerFactory;
  import org.apache.activemq.artemis.core.server.ActiveMQServer;
  import org.apache.activemq.artemis.core.server.cluster.ClusterConnection;
  import org.apache.activemq.artemis.core.server.cluster.MessageFlowRecord;
  import org.apache.activemq.artemis.core.server.cluster.impl.ClusterConnectionImpl;
  import org.apache.activemq.artemis.core.server.cluster.impl.MessageLoadBalancingType;
  import org.apache.activemq.artemis.core.settings.impl.AddressSettings;
  import org.apache.activemq.artemis.jms.client.ActiveMQConnectionFactory;
  import org.apache.activemq.artemis.protocol.amqp.broker.ProtonProtocolManagerFactory;
  import org.apache.activemq.artemis.spi.core.protocol.RemotingConnection;
  import org.apache.activemq.artemis.tests.extensions.parameterized.Parameter;
  import org.apache.activemq.artemis.tests.extensions.parameterized.ParameterizedTestExtension;
  import org.apache.activemq.artemis.tests.extensions.parameterized.Parameters;
  import org.apache.activemq.artemis.tests.integration.cluster.distribution.ClusterTestBase;
  import org.apache.activemq.artemis.tests.util.Wait;
  import org.apache.qpid.jms.JmsConnectionFactory;
  import org.junit.jupiter.api.BeforeEach;
  import org.junit.jupiter.api.TestTemplate;
  import org.junit.jupiter.api.extension.ExtendWith;
  import org.slf4j.Logger;
  import org.slf4j.LoggerFactory;
  
  @ExtendWith(ParameterizedTestExtension.class)
  public class ProtocolsMessageLoadBalancingTest extends ClusterTestBase {
  
     private static final Logger logger = LoggerFactory.getLogger(MethodHandles.lookup().lookupClass());
  
     private static final int NUMBER_OF_SERVERS = 2;
     private static final SimpleString queueName = SimpleString.of("queues.0");
  
     // I'm taking any number that /2 = Odd
     // to avoid perfect roundings and making sure messages are evenly distributed
     private static final int NUMBER_OF_MESSAGES = 77 * 2;
  
     @Parameters(name = "protocol={0}")
     public static Collection getParameters() {
        return Arrays.asList(new Object[][]{{"AMQP"}, {"CORE"}, {"OPENWIRE"}});
     }
  
     @Parameter(index = 0)
     public String protocol;
  
     @Override
     @BeforeEach
     public void setUp() throws Exception {
        super.setUp();
  
     }
  
     private void startServers(MessageLoadBalancingType loadBalancingType) throws Exception {
        setupServers();
  
        setRedistributionDelay(0);
  
        setupCluster(loadBalancingType);
  
        AddressSettings as = new AddressSettings().setRedistributionDelay(0).setExpiryAddress(SimpleString.of("queues.expiry"));
  
        getServer(0).getAddressSettingsRepository().addMatch("queues.*", as);
        getServer(1).getAddressSettingsRepository().addMatch("queues.*", as);
  
        startServers(0);
        startServers(1);
  
        createQueue(SimpleString.of("queues.expiry"));
        createQueue(queueName);
     }
  
     private void createQueue(SimpleString queueName) throws Exception {
        QueueConfiguration queueConfiguration = QueueConfiguration.of(queueName).setRoutingType(RoutingType.ANYCAST);
        servers[0].createQueue(queueConfiguration);
        servers[1].createQueue(queueConfiguration);
     }
  
     protected boolean isNetty() {
        return true;
     }
  
     private ConnectionFactory getJmsConnectionFactory(int node) {
        if (protocol.equals("AMQP")) {
           return new JmsConnectionFactory("amqp://localhost:" + (61616 + node));
        } else if (protocol.equals("OPENWIRE")) {
           return new org.apache.activemq.ActiveMQConnectionFactory("tcp://localhost:" + (61616 + node));
        } else if (protocol.equals("CORE")) {
           return new ActiveMQConnectionFactory("tcp://localhost:" + (61616 + node));
        } else {
           fail("Protocol " + protocol + " unknown");
           return null;
        }
     }
  
     private void pauseClusteringBridges(ActiveMQServer server) throws Exception {
        for (ClusterConnection clusterConnection : server.getClusterManager().getClusterConnections()) {
           for (MessageFlowRecord record : ((ClusterConnectionImpl) clusterConnection).getRecords().values()) {
              record.getBridge().pause();
           }
        }
     }
  
     @TestTemplate
     public void testLoadBalancing() throws Exception {
  
        startServers(MessageLoadBalancingType.STRICT);
  
        ConnectionFactory[] factory = new ConnectionFactory[NUMBER_OF_SERVERS];
        Connection[] connection = new Connection[NUMBER_OF_SERVERS];
        Session[] session = new Session[NUMBER_OF_SERVERS];
        MessageConsumer[] consumer = new MessageConsumer[NUMBER_OF_SERVERS];
  
        // this will pre create consumers to make sure messages are distributed evenly without redistribution
        for (int node = 0; node < NUMBER_OF_SERVERS; node++) {
           factory[node] = getJmsConnectionFactory(node);
           connection[node] = factory[node].createConnection();
           session[node] = connection[node].createSession(false, Session.AUTO_ACKNOWLEDGE);
           consumer[node] = session[node].createConsumer(session[node].createQueue(queueName.toString()));
        }
  
        waitForBindings(0, "queues.0", 1, 1, true);
        waitForBindings(1, "queues.0", 1, 1, true);
  
        waitForBindings(0, "queues.0", 1, 1, false);
        waitForBindings(1, "queues.0", 1, 1, false);
  
        pauseClusteringBridges(servers[0]);
  
        // sending Messages.. they should be load balanced
        {
           ConnectionFactory cf = getJmsConnectionFactory(0);
           Connection cn = cf.createConnection();
           runAfter(cn::close);
           Session sn = cn.createSession(false, Session.AUTO_ACKNOWLEDGE);
           MessageProducer pd = sn.createProducer(sn.createQueue(queueName.toString()));
  
           for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
              pd.send(sn.createTextMessage("hello " + i));
           }
  
           cn.close();
        }
  
        receiveMessages(connection[0], consumer[0], NUMBER_OF_MESSAGES / 2, true);
        connection[1].start();
        assertNull(consumer[1].receiveNoWait());
        connection[1].stop();
  
        servers[0].stop();
        clearServer(0);
  
        setupServer(0, isFileStorage(), isNetty());
        servers[0].addProtocolManagerFactory(new ProtonProtocolManagerFactory());
  
        setupClusterConnection("cluster0", "queues", MessageLoadBalancingType.STRICT, 1, isNetty(), 0, 1);
  
        servers[0].start();
  
        receiveMessages(connection[1], consumer[1], NUMBER_OF_MESSAGES / 2, true);
        for (int node = 0; node < NUMBER_OF_SERVERS; node++) {
           connection[node].close();
        }
  
     }
  
     @TestTemplate
     public void testRedistributionStoppedWithNoRemoteConsumers() throws Exception {
        startServers(MessageLoadBalancingType.ON_DEMAND);
  
        ConnectionFactory factory = getJmsConnectionFactory(1);
  
        // Wait for cluster nodes to sync
        waitForBindings(0, "queues.0", 1, 0, true);
        waitForBindings(1, "queues.0", 1, 0, true);
  
        waitForBindings(0, "queues.0", 1, 0, false);
        waitForBindings(1, "queues.0", 1, 0, false);
  
        // Create CFs
        ConnectionFactory cf0 = getJmsConnectionFactory(0);
        ConnectionFactory cf1 = getJmsConnectionFactory(1);
  
        // Create Consumers
        Connection cn0 = cf0.createConnection();
        runAfter(cn0::close);
        Session sn0 = cn0.createSession(false, Session.AUTO_ACKNOWLEDGE);
        MessageConsumer c0 = sn0.createConsumer(sn0.createQueue(queueName.toString()));
        cn0.start();
  
        Connection cn1 = cf1.createConnection();
        runAfter(cn1::close);
        Session sn1 = cn1.createSession(false, Session.AUTO_ACKNOWLEDGE);
        MessageConsumer c1 = sn1.createConsumer(sn0.createQueue(queueName.toString()));
        cn1.start();
  
        MessageProducer pd = sn0.createProducer(sn0.createQueue(queueName.toString()));
  
        // Wait for cluster nodes to sync consumer count
        waitForBindings(0, "queues.0", 1, 1, false);
        waitForBindings(1, "queues.0", 1, 1, false);
  
        // Start queue redistributor
        c0.close();
  
        // Send Messages to node1.
        for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
           pd.send(sn0.createTextMessage("hello " + i));
        }
  
        // Ensure the messages are redistributed from node0 to node1.
        for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
           assertNotNull(c1.receive(1000));
        }
  
        // Close client on node1.  This should make the node0 stop redistributing.
        c1.close();
        sn1.close();
        cn1.close();
  
        // Send more messages (these should stay in node0)
        for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
           pd.send(sn0.createTextMessage("hello " + i));
        }
  
        cn0.close();
  
        // Messages should stay in node 1 and note get redistributed.
        Wait.assertEquals(NUMBER_OF_MESSAGES, servers[0].locateQueue(queueName)::getMessageCount);
        Wait.assertEquals(0, servers[1].locateQueue(queueName)::getMessageCount);
     }
  
     @TestTemplate
     public void testExpireRedistributed() throws Exception {
        startServers(MessageLoadBalancingType.ON_DEMAND);
  
        ConnectionFactory factory = getJmsConnectionFactory(1);
  
        waitForBindings(0, "queues.0", 1, 0, true);
        waitForBindings(1, "queues.0", 1, 0, true);
  
        waitForBindings(0, "queues.0", 1, 0, false);
        waitForBindings(1, "queues.0", 1, 0, false);
  
        // sending Messages..
        {
           ConnectionFactory cf = getJmsConnectionFactory(0);
           Connection cn = cf.createConnection();
           Session sn = cn.createSession(false, Session.AUTO_ACKNOWLEDGE);
           MessageProducer pd = sn.createProducer(sn.createQueue(queueName.toString()));
           pd.setTimeToLive(200);
  
           for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
              pd.send(sn.createTextMessage("hello " + i));
           }
  
           cn.close();
        }
  
        // time to let stuff expire
        Thread.sleep(200);
  
        Connection connection = factory.createConnection();
        runAfter(connection::close);
        Session session = connection.createSession(false, Session.AUTO_ACKNOWLEDGE);
        MessageConsumer consumer = session.createConsumer(session.createQueue("queues.expiry"));
  
        receiveMessages(connection, consumer, NUMBER_OF_MESSAGES, true);
        connection.close();
     }
  
     @TestTemplate
     public void testRestartConnection() throws Exception {
  
        startServers(MessageLoadBalancingType.STRICT);
  
        logger.debug("connections {}", servers[1].getRemotingService().getConnections().size());
  
        Wait.assertEquals(3, () -> servers[1].getRemotingService().getConnections().size());
        Wait.assertEquals(3, () -> servers[0].getRemotingService().getConnections().size());
  
        RemotingConnection[] connectionsServer1 = servers[1].getRemotingService().getConnections().toArray(new RemotingConnection[3]);
        RemotingConnection[] connectionsServer0 = servers[0].getRemotingService().getConnections().toArray(new RemotingConnection[3]);
  
        ConnectionFactory[] factory = new ConnectionFactory[NUMBER_OF_SERVERS];
        Connection[] connection = new Connection[NUMBER_OF_SERVERS];
        Session[] session = new Session[NUMBER_OF_SERVERS];
        MessageConsumer[] consumer = new MessageConsumer[NUMBER_OF_SERVERS];
  
        // this will pre create consumers to make sure messages are distributed evenly without redistribution
        for (int node = 0; node < NUMBER_OF_SERVERS; node++) {
           factory[node] = getJmsConnectionFactory(node);
           connection[node] = factory[node].createConnection();
           runAfter(connection[node]::close);
           session[node] = connection[node].createSession(false, Session.AUTO_ACKNOWLEDGE);
           consumer[node] = session[node].createConsumer(session[node].createQueue(queueName.toString()));
        }
  
        waitForBindings(0, "queues.0", 1, 1, true);
        waitForBindings(1, "queues.0", 1, 1, true);
  
        waitForBindings(0, "queues.0", 1, 1, false);
        waitForBindings(1, "queues.0", 1, 1, false);
  
        for (RemotingConnection remotingConnection : servers[1].getRemotingService().getConnections()) {
           remotingConnection.fail(new ActiveMQException("forcing failure"));
        }
        for (RemotingConnection remotingConnection : servers[1].getRemotingService().getConnections()) {
           remotingConnection.fail(new ActiveMQException("forcing failure"));
        }
  
        // this is to allow reconnects
        Thread.sleep(500);
  
        // this will pre create consumers to make sure messages are distributed evenly without redistribution
        for (int node = 0; node < NUMBER_OF_SERVERS; node++) {
           try {
              connection[node].close();
           } catch (Throwable e) {
              e.printStackTrace();
           }
           factory[node] = getJmsConnectionFactory(node);
           connection[node] = factory[node].createConnection();
           session[node] = connection[node].createSession(false, Session.AUTO_ACKNOWLEDGE);
           consumer[node] = session[node].createConsumer(session[node].createQueue(queueName.toString()));
        }
  
        waitForBindings(0, "queues.0", 1, 1, true);
        waitForBindings(1, "queues.0", 1, 1, true);
  
        waitForBindings(0, "queues.0", 1, 1, false);
        waitForBindings(1, "queues.0", 1, 1, false);
  
        logger.debug("connections {}", servers[1].getRemotingService().getConnections().size());
  
        // sending Messages.. they should be load balanced
        {
           ConnectionFactory cf = getJmsConnectionFactory(0);
           Connection cn = cf.createConnection();
           Session sn = cn.createSession(false, Session.AUTO_ACKNOWLEDGE);
           MessageProducer pd = sn.createProducer(sn.createQueue(queueName.toString()));
  
           for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
              pd.send(sn.createTextMessage("hello " + i));
           }
  
           cn.close();
        }
  
        receiveMessages(connection[0], consumer[0], NUMBER_OF_MESSAGES / 2, true);
        receiveMessages(connection[1], consumer[1], NUMBER_OF_MESSAGES / 2, true);
  
        for (int node = 0; node < NUMBER_OF_SERVERS; node++) {
           connection[node].close();
        }
  
     }
  
     @TestTemplate
     public void testRedistributeAfterLoadBalanced() throws Exception {
  
        startServers(MessageLoadBalancingType.ON_DEMAND);
  
        ConnectionFactory[] factory = new ConnectionFactory[NUMBER_OF_SERVERS];
        Connection[] connection = new Connection[NUMBER_OF_SERVERS];
        Session[]  session = new Session[NUMBER_OF_SERVERS];
        MessageConsumer[]  consumer = new MessageConsumer[NUMBER_OF_SERVERS];
  
        // this will pre create consumers to make sure messages are distributed evenly without redistribution
        for (int node = 0; node < NUMBER_OF_SERVERS; node++) {
           factory[node] = getJmsConnectionFactory(node);
           connection[node] = factory[node].createConnection();
           runAfter(connection[node]::close);
           session[node] = connection[node].createSession(false, Session.AUTO_ACKNOWLEDGE);
           consumer[node] = session[node].createConsumer(session[node].createQueue(queueName.toString()));
        }
  
        waitForBindings(0, "queues.0", 1, 1, true);
        waitForBindings(1, "queues.0", 1, 1, true);
  
        waitForBindings(0, "queues.0", 1, 1, false);
        waitForBindings(1, "queues.0", 1, 1, false);
  
  
  
        if (protocol.equals("AMQP")) {
           ServerLocator locator = ActiveMQClient.createServerLocator("tcp://localhost:61616");
           locator.setMinLargeMessageSize(1024);
           ClientSessionFactory coreFactory = locator.createSessionFactory();
           ClientSession clientSession = coreFactory.createSession();
           ClientProducer producer = clientSession.createProducer(queueName);
           String largeBody = " ".repeat(10 * 1024);
           for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
              ClientMessage message = clientSession.createMessage((byte)0, true);
              // making 1/3 of the messages to be large
              message.getBodyBuffer().writeUTF(i % 3 == 0 ? largeBody : "hello");
              producer.send(message);
           }
           coreFactory.close();
  
        } else {
  
           // sending Messages.. they should be load balanced
           ConnectionFactory cf = getJmsConnectionFactory(0);
           Connection cn = cf.createConnection();
           Session sn = cn.createSession(false, Session.AUTO_ACKNOWLEDGE);
           MessageProducer pd = sn.createProducer(sn.createQueue(queueName.toString()));
  
           for (int i = 0; i < NUMBER_OF_MESSAGES; i++) {
              StringBuilder sb = new StringBuilder();
              sb.append("hello");
              Message message = sn.createTextMessage(sb.toString());
              pd.send(message);
           }
  
           cn.close();
        }
  
        Wait.assertEquals(NUMBER_OF_MESSAGES / 2, servers[1].locateQueue(queueName)::getMessageCount);
        Wait.assertEquals(NUMBER_OF_MESSAGES / 2, servers[0].locateQueue(queueName)::getMessageCount);
        receiveMessages(connection[0], consumer[0], NUMBER_OF_MESSAGES / 2, true);
        connection[1].close();
        // this wil be after redistribution
        receiveMessages(connection[0], consumer[0], NUMBER_OF_MESSAGES / 2, true);
        connection[0].close();
     }
  
  
     private void receiveMessages(Connection connection,
                                  MessageConsumer messageConsumer,
                                  int messageCount,
                                  boolean exactCount) throws JMSException {
        connection.start();
  
        for (int i = 0; i < messageCount; i++) {
           Message msg = messageConsumer.receive(5000);
           assertNotNull(msg, "did not receive message at " + i);
        }
  
        // this means no more messages received
        if (exactCount) {
           assertNull(messageConsumer.receiveNoWait());
        }
     }
  
     protected void setupCluster(final MessageLoadBalancingType messageLoadBalancingType) throws Exception {
        setupClusterConnection("cluster0", "queues", messageLoadBalancingType, 1, isNetty(), 0, 1);
  
        setupClusterConnection("cluster1", "queues", messageLoadBalancingType, 1, isNetty(), 1, 0);
     }
  
     protected void setRedistributionDelay(final long delay) {
     }
  
     protected void setupServers() throws Exception {
        setupServer(0, isFileStorage(), isNetty());
        setupServer(1, isFileStorage(), isNetty());
  
        servers[0].addProtocolManagerFactory(new ProtonProtocolManagerFactory());
        servers[1].addProtocolManagerFactory(new ProtonProtocolManagerFactory());
        servers[0].addProtocolManagerFactory(new OpenWireProtocolManagerFactory());
        servers[1].addProtocolManagerFactory(new OpenWireProtocolManagerFactory());
     }
  
     protected void stopServers() throws Exception {
        closeAllConsumers();
  
        closeAllSessionFactories();
  
        closeAllServerLocatorsFactories();
  
        stopServers(0, 1);
  
        clearServer(0, 1);
     }
  
     @Override
     protected ConfigurationImpl createBasicConfig(final int serverID) {
        ConfigurationImpl configuration = super.createBasicConfig(serverID);
        configuration.setMessageExpiryScanPeriod(100);
  
        return configuration;
     }
  
  }
  