/**
* 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.tool;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Iterator;
import java.util.List;
import java.util.Random;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import javax.jms.BytesMessage;
import javax.jms.Connection;
import javax.jms.DeliveryMode;
import javax.jms.Destination;
import javax.jms.JMSException;
import javax.jms.Message;
import javax.jms.MessageConsumer;
import javax.jms.MessageProducer;
import javax.jms.Session;
import junit.framework.TestCase;
import org.apache.activemq.ActiveMQConnectionFactory;
import org.apache.activemq.command.ActiveMQQueue;
/**
*
*/
public class AcidTestTool extends TestCase {
// Worker configuration.
protected int recordSize = 1024;
protected int batchSize = 5;
protected int workerThinkTime = 500;
protected Destination target;
private Random random = new Random();
private byte data[];
private int workerCount = 10;
private AtomicBoolean ignoreJMSErrors = new AtomicBoolean(false);
private ActiveMQConnectionFactory factory;
private Connection connection;
private AtomicInteger publishedBatches = new AtomicInteger(0);
private AtomicInteger consumedBatches = new AtomicInteger(0);
private List<Throwable> errors = Collections.synchronizedList(new ArrayList<Throwable>());
private interface Worker extends Runnable {
boolean waitForExit(long i) throws InterruptedException;
}
private final class ProducerWorker implements Worker {
private Session session;
private MessageProducer producer;
private BytesMessage message;
private CountDownLatch doneLatch = new CountDownLatch(1);
ProducerWorker(Session session, String workerId) throws JMSException {
this.session = session;
producer = session.createProducer(target);
producer.setDeliveryMode(DeliveryMode.PERSISTENT);
message = session.createBytesMessage();
message.setStringProperty("workerId", workerId);
message.writeBytes(data);
}
public void run() {
try {
for (int batchId = 0; true; batchId++) {
// System.out.println("Sending batch: "+workerId+"
// "+batchId);
for (int msgId = 0; msgId < batchSize; msgId++) {
// Sleep some random amount of time less than
// workerThinkTime
try {
Thread.sleep(random.nextInt(workerThinkTime));
} catch (InterruptedException e1) {
return;
}
message.setIntProperty("batch-id", batchId);
message.setIntProperty("msg-id", msgId);
producer.send(message);
}
session.commit();
publishedBatches.incrementAndGet();
// System.out.println("Commited send batch: "+workerId+"
// "+batchId);
}
} catch (JMSException e) {
if (!ignoreJMSErrors.get()) {
e.printStackTrace();
errors.add(e);
}
return;
} catch (Throwable e) {
e.printStackTrace();
errors.add(e);
return;
} finally {
System.out.println("Producer exiting.");
doneLatch.countDown();
}
}
public boolean waitForExit(long i) throws InterruptedException {
return doneLatch.await(i, TimeUnit.MILLISECONDS);
}
}
private final class ConsumerWorker implements Worker {
private Session session;
private MessageConsumer consumer;
private final long timeout;
private CountDownLatch doneLatch = new CountDownLatch(1);
ConsumerWorker(Session session, String workerId, long timeout) throws JMSException {
this.session = session;
this.timeout = timeout;
consumer = session.createConsumer(target, "workerId='" + workerId + "'");
}
public void run() {
try {
int batchId = 0;
while (true) {
for (int msgId = 0; msgId < batchSize; msgId++) {
// Sleep some random amount of time less than
// workerThinkTime
try {
Thread.sleep(random.nextInt(workerThinkTime));
} catch (InterruptedException e1) {
return;
}
Message message = consumer.receive(timeout);
if (msgId > 0) {
assertNotNull(message);
assertEquals(message.getIntProperty("batch-id"), batchId);
assertEquals(message.getIntProperty("msg-id"), msgId);
} else {
if (message == null) {
System.out.println("At end of batch an don't have a next batch to process. done.");
return;
}
assertEquals(msgId, message.getIntProperty("msg-id"));
batchId = message.getIntProperty("batch-id");
// System.out.println("Receiving batch: "+workerId+"
// "+batchId);
}
}
session.commit();
consumedBatches.incrementAndGet();
// System.out.println("Commited receive batch: "+workerId+"
// "+batchId);
}
} catch (JMSException e) {
if (!ignoreJMSErrors.get()) {
e.printStackTrace();
errors.add(e);
}
return;
} catch (Throwable e) {
e.printStackTrace();
errors.add(e);
return;
} finally {
System.out.println("Consumer exiting.");
doneLatch.countDown();
}
}
public boolean waitForExit(long i) throws InterruptedException {
return doneLatch.await(i, TimeUnit.MILLISECONDS);
}
}
/**
* @see junit.framework.TestCase#setUp()
*/
protected void setUp() throws Exception {
factory = new ActiveMQConnectionFactory("tcp://localhost:61616");
this.target = new ActiveMQQueue(getClass().getName());
}
protected void tearDown() throws Exception {
if (connection != null) {
try {
connection.close();
} catch (Throwable ignore) {
}
connection = null;
}
}
/**
* @throws InterruptedException
* @throws JMSException
* @throws JMSException
*/
private void reconnect() throws InterruptedException, JMSException {
if (connection != null) {
try {
connection.close();
} catch (Throwable ignore) {
}
connection = null;
}
long reconnectDelay = 1000;
while (connection == null) {
if (reconnectDelay > 1000 * 10) {
reconnectDelay = 1000 * 10;
}
try {
connection = factory.createConnection();
connection.start();
} catch (JMSException e) {
Thread.sleep(reconnectDelay);
reconnectDelay *= 2;
}
}
}
/**
* @throws Throwable
* @throws IOException
*/
public void testAcidTransactions() throws Throwable {
System.out.println("Client threads write records using: Record Size: " + recordSize + ", Batch Size: " + batchSize + ", Worker Think Time: " + workerThinkTime);
// Create the record and fill it with some values.
data = new byte[recordSize];
for (int i = 0; i < data.length; i++) {
data[i] = (byte)i;
}
System.out.println("==============================================");
System.out.println("===> Start the server now.");
System.out.println("==============================================");
reconnect();
System.out.println("Starting " + workerCount + " Workers...");
ArrayList<Worker> workers = new ArrayList<Worker>();
for (int i = 0; i < workerCount; i++) {
String workerId = "worker-" + i;
Worker w = new ConsumerWorker(connection.createSession(true, Session.SESSION_TRANSACTED), workerId, 1000 * 5);
workers.add(w);
new Thread(w, "Consumer:" + workerId).start();
w = new ProducerWorker(connection.createSession(true, Session.SESSION_TRANSACTED), workerId);
workers.add(w);
new Thread(w, "Producer:" + workerId).start();
}
System.out.println("Waiting for " + (workerCount * 10) + " batches to be delivered.");
//
// Wait for about 5 batches of messages per worker to be consumed before
// restart.
//
while (publishedBatches.get() < workerCount * 5) {
System.out.println("Stats: Produced Batches: " + this.publishedBatches.get() + ", Consumed Batches: " + this.consumedBatches.get());
Thread.sleep(1000);
}
System.out.println("==============================================");
System.out.println("===> Server is under load now. Kill it!");
System.out.println("==============================================");
ignoreJMSErrors.set(true);
// Wait for all the workers to finish.
System.out.println("Waiting for all workers to exit due to server shutdown.");
for (Iterator<Worker> iter = workers.iterator(); iter.hasNext();) {
Worker worker = iter.next();
while (!worker.waitForExit(1000)) {
System.out.println("==============================================");
System.out.println("===> Server is under load now. Kill it!");
System.out.println("==============================================");
System.out.println("Stats: Produced Batches: " + this.publishedBatches.get() + ", Consumed Batches: " + this.consumedBatches.get());
}
}
workers.clear();
// No errors should have occurred so far.
if (errors.size() > 0) {
throw errors.get(0);
}
System.out.println("==============================================");
System.out.println("===> Start the server now.");
System.out.println("==============================================");
reconnect();
System.out.println("Restarted.");
// Validate the all transactions were commited as a uow. Looking for
// partial commits.
for (int i = 0; i < workerCount; i++) {
String workerId = "worker-" + i;
Worker w = new ConsumerWorker(connection.createSession(true, Session.SESSION_TRANSACTED), workerId, 5 * 1000);
workers.add(w);
new Thread(w, "Consumer:" + workerId).start();
}
System.out.println("Waiting for restarted consumers to finish consuming all messages..");
for (Iterator<Worker> iter = workers.iterator(); iter.hasNext();) {
Worker worker = iter.next();
while (!worker.waitForExit(1000 * 5)) {
System.out.println("Waiting for restarted consumers to finish consuming all messages..");
System.out.println("Stats: Produced Batches: " + this.publishedBatches.get() + ", Consumed Batches: " + this.consumedBatches.get());
}
}
workers.clear();
System.out.println("Workers finished..");
System.out.println("Stats: Produced Batches: " + this.publishedBatches.get() + ", Consumed Batches: " + this.consumedBatches.get());
if (errors.size() > 0) {
throw errors.get(0);
}
}
public static void main(String[] args) {
try {
AcidTestTool tool = new AcidTestTool();
tool.setUp();
tool.testAcidTransactions();
tool.tearDown();
} catch (Throwable e) {
System.out.println("Test Failed: " + e.getMessage());
e.printStackTrace();
}
}
}