Merge pull request #24 from rabbitmq/rabbitmq-perf-test-23

Provide a way to load message payload from disk

Author: michaelklishin

src/main/java/com/rabbitmq/perf/LocalFilesMessageBodySource.java

@@ -0,0 +1,64 @@
+// Copyright (c) 2017-Present Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
+// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
+// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+
+package com.rabbitmq.perf;
+
+import java.io.*;
+import java.util.ArrayList;
+import java.util.List;
+
+/**
+ *
+ */
+public class LocalFilesMessageBodySource implements MessageBodySource {
+
+    private final List<byte[]> bodies;
+
+    private final String contentType;
+
+    public LocalFilesMessageBodySource(List<String> filesNames, String contentType) throws IOException {
+        bodies = new ArrayList<byte[]>(filesNames.size());
+        for (String fileName : filesNames) {
+            File file = new File(fileName.trim());
+            if (!file.exists() || file.isDirectory()) {
+                throw new IllegalArgumentException(fileName + " isn't a valid body file.");
+            }
+            BufferedInputStream inputStream = null;
+            try {
+                inputStream = new BufferedInputStream(new FileInputStream(file));
+                byte [] body = new byte[(int) file.length()];
+                inputStream.read(body, 0, body.length);
+                bodies.add(body);
+            } finally {
+                if (inputStream != null) {
+                    inputStream.close();
+                }
+            }
+
+        }
+        this.contentType = contentType;
+    }
+
+    public LocalFilesMessageBodySource(List<String> filesNames) throws IOException {
+        this(filesNames, null);
+    }
+
+    @Override
+    public MessageBodyAndContentType create(int sequenceNumber) throws IOException {
+        return new MessageBodyAndContentType(
+            bodies.get(sequenceNumber % bodies.size()), contentType
+        );
+    }
+}

src/main/java/com/rabbitmq/perf/MessageBodySource.java

@@ -0,0 +1,46 @@
+// Copyright (c) 2017-Present Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
+// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
+// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+
+package com.rabbitmq.perf;
+
+import java.io.IOException;
+
+/**
+ * Sources produce message bodies and content type
+ * used by publishers.
+ */
+public interface MessageBodySource {
+
+    MessageBodyAndContentType create(int sequenceNumber) throws IOException;
+
+    class MessageBodyAndContentType {
+        private final byte [] body;
+        private final String contentType;
+
+        public MessageBodyAndContentType(byte[] body, String contentType) {
+            this.body = body;
+            this.contentType = contentType;
+        }
+
+        public byte[] getBody() {
+            return body;
+        }
+
+        public String getContentType() {
+            return contentType;
+        }
+    }
+
+}

src/main/java/com/rabbitmq/perf/MulticastParams.java

@@ -54,6 +54,9 @@ public class MulticastParams {
     private boolean autoAck = true;
     private boolean autoDelete = false;
 
+    private List<String> bodyFiles = new ArrayList<String>();
+    private String bodyContentType = null;
+
     private boolean predeclared;
 
     public void setExchangeType(String exchangeType) {
@@ -201,18 +204,36 @@ public boolean getRandomRoutingKey() {
         return randomRoutingKey;
     }
 
+    public void setBodyFiles(List<String> bodyFiles) {
+        if (bodyFiles == null) {
+            this.bodyFiles = new ArrayList<String>();
+        } else {
+            this.bodyFiles = new ArrayList<String>(bodyFiles);
+        }
+    }
+
+    public void setBodyContentType(String bodyContentType) {
+        this.bodyContentType = bodyContentType;
+    }
+
     public Producer createProducer(Connection connection, Stats stats, String id) throws IOException {
         Channel channel = connection.createChannel();
         if (producerTxSize > 0) channel.txSelect();
         if (confirm >= 0) channel.confirmSelect();
         if (!predeclared || !exchangeExists(connection, exchangeName)) {
             channel.exchangeDeclare(exchangeName, exchangeType);
         }
+        MessageBodySource messageBodySource = null;
+        if (bodyFiles.size() > 0) {
+            messageBodySource = new LocalFilesMessageBodySource(bodyFiles, bodyContentType);
+        } else {
+            messageBodySource = new TimeSequenceMessageBodySource(minMsgSize);
+        }
         final Producer producer = new Producer(channel, exchangeName, id,
                                                randomRoutingKey, flags, producerTxSize,
                                                producerRateLimit, producerMsgCount,
-                                               minMsgSize, timeLimit,
-                                               confirm, stats);
+                                               timeLimit,
+                                               confirm, messageBodySource, stats);
         channel.addReturnListener(producer);
         channel.addConfirmListener(producer);
         return producer;

src/main/java/com/rabbitmq/perf/PerfTest.java

@@ -31,6 +31,7 @@
 
 import com.rabbitmq.client.ConnectionFactory;
 
+import static java.util.Arrays.asList;
 
 public class PerfTest {
 
@@ -73,6 +74,8 @@ public static void main(String[] args) {
             List<?> flags            = lstArg(cmd, 'f');
             int frameMax             = intArg(cmd, 'M', 0);
             int heartbeat            = intArg(cmd, 'b', 0);
+            String bodyFiles         = strArg(cmd, 'B', null);
+            String bodyContentType   = strArg(cmd, 'T', null);
             boolean predeclared      = cmd.hasOption('p');
 
             String uri               = strArg(cmd, 'h', "amqp://localhost");
@@ -83,7 +86,7 @@ public static void main(String[] args) {
                 for(int i = 0; i< urisArray.length; i++) {
                     urisArray[i] = urisArray[i].trim();
                 }
-                uris = Arrays.asList(urisArray);
+                uris = asList(urisArray);
             } else {
                 uris = Collections.singletonList(uri);
             }
@@ -124,11 +127,13 @@ public static void main(String[] args) {
             p.setProducerChannelCount(  producerChannelCount);
             p.setProducerMsgCount(      producerMsgCount);
             p.setProducerTxSize(        producerTxSize);
-            p.setQueueNames(            Arrays.asList(queueNames.split(",")));
+            p.setQueueNames(            asList(queueNames.split(",")));
             p.setRoutingKey(            routingKey);
             p.setRandomRoutingKey(      randomRoutingKey);
             p.setProducerRateLimit(     producerRateLimit);
             p.setTimeLimit(             timeLimit);
+            p.setBodyFiles(             bodyFiles == null ? null : asList(bodyFiles.split(",")));
+            p.setBodyContentType(       bodyContentType);
 
             MulticastSet set = new MulticastSet(stats, factory, p, testID, uris);
             set.run(true);
@@ -185,6 +190,9 @@ private static Options getOptions() {
         options.addOption(new Option("M", "framemax",               true, "frame max"));
         options.addOption(new Option("b", "heartbeat",              true, "heartbeat interval"));
         options.addOption(new Option("p", "predeclared",            false,"allow use of predeclared objects"));
+        options.addOption(new Option("B", "body",                   true, "comma-separated list of files to use in message bodies"));
+        options.addOption(new Option("T", "bodyContenType",         true, "body content-type"));
+
         return options;
     }
 
@@ -205,7 +213,7 @@ private static List<?> lstArg(CommandLine cmd, char opt) {
         if (vals == null) {
             vals = new String[] {};
         }
-        return Arrays.asList(vals);
+        return asList(vals);
     }
 
     private static class PrintlnStats extends Stats {

src/main/java/com/rabbitmq/perf/Producer.java

@@ -18,11 +18,8 @@
 import com.rabbitmq.client.AMQP;
 import com.rabbitmq.client.Channel;
 import com.rabbitmq.client.ConfirmListener;
-import com.rabbitmq.client.MessageProperties;
 import com.rabbitmq.client.ReturnListener;
 
-import java.io.ByteArrayOutputStream;
-import java.io.DataOutputStream;
 import java.io.IOException;
 import java.util.Collections;
 import java.util.List;
@@ -47,30 +44,30 @@ public class Producer extends ProducerConsumerBase implements Runnable, ReturnLi
 
     private final Stats   stats;
 
-    private final byte[]  message;
+    private final MessageBodySource messageBodySource;
 
     private Semaphore confirmPool;
     private final SortedSet<Long> unconfirmedSet =
         Collections.synchronizedSortedSet(new TreeSet<Long>());
 
     public Producer(Channel channel, String exchangeName, String id, boolean randomRoutingKey,
                     List<?> flags, int txSize,
-                    float rateLimit, int msgLimit, int minMsgSize, int timeLimit,
-                    long confirm, Stats stats)
+                    float rateLimit, int msgLimit, int timeLimit,
+                    long confirm, MessageBodySource messageBodySource, Stats stats)
         throws IOException {
 
-        this.channel          = channel;
-        this.exchangeName     = exchangeName;
-        this.id               = id;
-        this.randomRoutingKey = randomRoutingKey;
-        this.mandatory        = flags.contains("mandatory");
-        this.immediate        = flags.contains("immediate");
-        this.persistent       = flags.contains("persistent");
-        this.txSize           = txSize;
-        this.rateLimit        = rateLimit;
-        this.msgLimit         = msgLimit;
-        this.timeLimit        = 1000L * timeLimit;
-        this.message          = new byte[minMsgSize];
+        this.channel            = channel;
+        this.exchangeName       = exchangeName;
+        this.id                 = id;
+        this.randomRoutingKey   = randomRoutingKey;
+        this.mandatory          = flags.contains("mandatory");
+        this.immediate          = flags.contains("immediate");
+        this.persistent         = flags.contains("persistent");
+        this.txSize             = txSize;
+        this.rateLimit          = rateLimit;
+        this.msgLimit           = msgLimit;
+        this.timeLimit          = 1000L * timeLimit;
+        this.messageBodySource = messageBodySource;
         if (confirm > 0) {
             this.confirmPool  = new Semaphore((int)confirm);
         }
@@ -136,7 +133,7 @@ public void run() {
                 if (confirmPool != null) {
                     confirmPool.acquire();
                 }
-                publish(createMessage(totalMsgCount));
+                publish(messageBodySource.create(totalMsgCount));
                 totalMsgCount++;
                 msgCount++;
 
@@ -154,33 +151,23 @@ public void run() {
         }
     }
 
-    private void publish(byte[] msg)
+    private void publish(MessageBodySource.MessageBodyAndContentType messageBodyAndContentType)
         throws IOException {
 
+        AMQP.BasicProperties.Builder propertiesBuilder = new AMQP.BasicProperties.Builder();
+        if (persistent) {
+            propertiesBuilder.deliveryMode(2);
+        }
+
+        if (messageBodyAndContentType.getContentType() != null) {
+            propertiesBuilder.contentType(messageBodyAndContentType.getContentType());
+        }
+
         unconfirmedSet.add(channel.getNextPublishSeqNo());
         channel.basicPublish(exchangeName, randomRoutingKey ? UUID.randomUUID().toString() : id,
                              mandatory, immediate,
-                             persistent ? MessageProperties.MINIMAL_PERSISTENT_BASIC : MessageProperties.MINIMAL_BASIC,
-                             msg);
-    }
-
-    private byte[] createMessage(int sequenceNumber)
-        throws IOException {
-
-        ByteArrayOutputStream acc = new ByteArrayOutputStream();
-        DataOutputStream d = new DataOutputStream(acc);
-        long nano = System.nanoTime();
-        d.writeInt(sequenceNumber);
-        d.writeLong(nano);
-        d.flush();
-        acc.flush();
-        byte[] m = acc.toByteArray();
-        if (m.length <= message.length) {
-            System.arraycopy(m, 0, message, 0, m.length);
-            return message;
-        } else {
-            return m;
-        }
+                             propertiesBuilder.build(),
+                             messageBodyAndContentType.getBody());
     }
 
 }

src/main/java/com/rabbitmq/perf/TimeSequenceMessageBodySource.java

@@ -0,0 +1,50 @@
+// Copyright (c) 2017-Present Pivotal Software, Inc.  All rights reserved.
+//
+// This software, the RabbitMQ Java client library, is triple-licensed under the
+// Mozilla Public License 1.1 ("MPL"), the GNU General Public License version 2
+// ("GPL") and the Apache License version 2 ("ASL"). For the MPL, please see
+// LICENSE-MPL-RabbitMQ. For the GPL, please see LICENSE-GPL2.  For the ASL,
+// please see LICENSE-APACHE2.
+//
+// This software is distributed on an "AS IS" basis, WITHOUT WARRANTY OF ANY KIND,
+// either express or implied. See the LICENSE file for specific language governing
+// rights and limitations of this software.
+//
+// If you have any questions regarding licensing, please contact us at
+// [email protected].
+
+package com.rabbitmq.perf;
+
+import java.io.ByteArrayOutputStream;
+import java.io.DataOutputStream;
+import java.io.IOException;
+
+/**
+ *
+ */
+public class TimeSequenceMessageBodySource implements MessageBodySource {
+
+    private final byte[] message;
+
+    public TimeSequenceMessageBodySource(int minMsgSize) {
+        this.message = new byte[minMsgSize];
+    }
+
+    @Override
+    public MessageBodyAndContentType create(int sequenceNumber) throws IOException {
+        ByteArrayOutputStream acc = new ByteArrayOutputStream();
+        DataOutputStream d = new DataOutputStream(acc);
+        long nano = System.nanoTime();
+        d.writeInt(sequenceNumber);
+        d.writeLong(nano);
+        d.flush();
+        acc.flush();
+        byte[] m = acc.toByteArray();
+        if (m.length <= message.length) {
+            System.arraycopy(m, 0, message, 0, m.length);
+            return new MessageBodyAndContentType(message, null);
+        } else {
+            return new MessageBodyAndContentType(m, null);
+        }
+    }
+}