com.dxc.temp.SimpleProducerConsumer.java Source code

Java tutorial

Introduction

Here is the source code for com.dxc.temp.SimpleProducerConsumer.java

Source

/*
 * Copyright 2010-2012 Amazon.com, Inc. or its affiliates. All Rights Reserved.
 *
 * Licensed under the Apache License, Version 2.0 (the "License").
 * You may not use this file except in compliance with the License.
 * A copy of the License is located at
 *
 *  http://aws.amazon.com/apache2.0
 *
 * or in the "license" file accompanying this file. This file 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 com.dxc.temp;

import java.math.BigInteger;
import java.util.ArrayList;
import java.util.List;
import java.util.Random;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;

import org.apache.commons.logging.Log;
import org.apache.commons.logging.LogFactory;

import com.amazonaws.AmazonClientException;
import com.amazonaws.ClientConfiguration;
import com.amazonaws.auth.AWSCredentials;
import com.amazonaws.auth.BasicAWSCredentials;
import com.amazonaws.services.sqs.AmazonSQS;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.amazonaws.services.sqs.model.BatchResultErrorEntry;
import com.amazonaws.services.sqs.model.DeleteMessageBatchRequest;
import com.amazonaws.services.sqs.model.DeleteMessageBatchRequestEntry;
import com.amazonaws.services.sqs.model.DeleteMessageBatchResult;
import com.amazonaws.services.sqs.model.DeleteMessageRequest;
import com.amazonaws.services.sqs.model.GetQueueUrlRequest;
import com.amazonaws.services.sqs.model.Message;
import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
import com.amazonaws.services.sqs.model.ReceiveMessageResult;
import com.amazonaws.services.sqs.model.SendMessageBatchRequest;
import com.amazonaws.services.sqs.model.SendMessageBatchRequestEntry;
import com.amazonaws.services.sqs.model.SendMessageBatchResult;
import com.amazonaws.services.sqs.model.SendMessageRequest;

/**
 * Start a specified number of producer and consumer threads, and
 * produce-consume for the least of the specified duration and 1h. Some messages
 * can be left in the queue because producers and consumers may not be in exact
 * balance.
 * <p>
 * The program does not validate its arguments.
 * <p>
 * An example command line (omitting class path for clarity) to produce-consume
 * with 1 producer and 2 consumers, batches of 10, for 20min is as follows:
 * <p>
 * {@code java com.amazonaws.sqs.samples.SimpleProducerConsumer [accessKey]
 * [secretKey] https://sqs.us-east-1.amazonaws.com 1 2 10 1024 20}
 * <p>
 * The command line parameters are, in order:
 *
 * @param accessKey
 *          The AWS access key
 * @param secretKey
 *          The AWS secret key
 * @param endpoint
 *          The SQS region endpoint
 * @param queueName
 *          The name of the queue. The program assumes that the queue already
 *          exists.
 * @param producerCount
 *          The number of producer threads to use
 * @param consumerCount
 *          The number of consumer threads to use
 * @param batchSize
 *          The size of batches to send, receive and delete messages. Must be
 *          between 1 and 10. With a value of 1, the program uses the single
 *          operation APIs ({@code SendMessage} and {@code DeleteMessage}).
 * @param messageSizeByte
 *          The size of messages to send in bytes. Note that the maximum batch
 *          size is 64KB so that batchSize * messageSizeByte must be < 64KB.
 * @param runTimeMinutes
 *          The time to run the program for. The program will run for the least
 *          of this duration and 1h.
 */
public class SimpleProducerConsumer {
    private static Log log = LogFactory.getLog(SimpleProducerConsumer.class);

    // maximum runtime of the program
    private static int MAX_RUNTIME_MINUTES = 60;

    public static void main(String[] args) throws InterruptedException {
        int argIndex = 0;

        final String accessKey = args[argIndex++];
        final String secretKey = args[argIndex++];
        final AWSCredentials credentials = new BasicAWSCredentials(accessKey, secretKey);

        final String endpoint = args[argIndex++];
        final String queueName = args[argIndex++];
        final int producerCount = Integer.parseInt(args[argIndex++]);
        final int consumerCount = Integer.parseInt(args[argIndex++]);
        final int batchSize = Integer.parseInt(args[argIndex++]);
        final int messageSizeByte = Integer.parseInt(args[argIndex++]);
        final int runTimeMinutes = Integer.parseInt(args[argIndex++]);

        // configure the SQS client with enough connections for all producer and
        // consumer threads
        AmazonSQS sqsClient = new AmazonSQSClient(credentials,
                new ClientConfiguration().withMaxConnections(producerCount + consumerCount));
        sqsClient.setEndpoint(endpoint);
        String queueUrl = sqsClient.getQueueUrl(new GetQueueUrlRequest(queueName)).getQueueUrl();

        // the flag to stop producer, consumer, and monitor threads
        AtomicBoolean stop = new AtomicBoolean(false);

        // start the producers
        final AtomicInteger producedCount = new AtomicInteger();
        Thread[] producers = new Thread[producerCount];
        for (int i = 0; i < producerCount; i++) {
            if (batchSize == 1)
                producers[i] = new Producer(sqsClient, queueUrl, messageSizeByte, producedCount, stop);
            else
                producers[i] = new BatchProducer(sqsClient, queueUrl, batchSize, messageSizeByte, producedCount,
                        stop);
            producers[i].start();
        }

        // start the consumers
        final AtomicInteger consumedCount = new AtomicInteger();
        Thread[] consumers = new Thread[consumerCount];
        for (int i = 0; i < consumerCount; i++) {
            if (batchSize == 1)
                consumers[i] = new Consumer(sqsClient, queueUrl, consumedCount, stop);
            else
                consumers[i] = new BatchConsumer(sqsClient, queueUrl, batchSize, consumedCount, stop);
            consumers[i].start();
        }

        // start the monitor (thread)
        Thread monitor = new Monitor(producedCount, consumedCount, stop);
        monitor.start();

        // wait for the specified amount of time then stop
        Thread.sleep(TimeUnit.MINUTES.toMillis(Math.min(runTimeMinutes, MAX_RUNTIME_MINUTES)));
        stop.set(true);

        // join all threads
        for (int i = 0; i < producerCount; i++)
            producers[i].join();

        for (int i = 0; i < consumerCount; i++)
            consumers[i].join();

        monitor.interrupt();
        monitor.join();
    }

    /**
     * Producer thread using {@code SendMessage} to send messages until stopped.
     */
    private static class Producer extends Thread {
        final AmazonSQS sqsClient;
        final String queueUrl;
        final AtomicInteger producedCount;
        final AtomicBoolean stop;
        final String theMessage;

        Producer(AmazonSQS sqsQueueBuffer, String queueUrl, int messageSizeByte, AtomicInteger producedCount,
                AtomicBoolean stop) {
            this.sqsClient = sqsQueueBuffer;
            this.queueUrl = queueUrl;
            this.producedCount = producedCount;
            this.stop = stop;
            this.theMessage = makeRandomString(messageSizeByte);
        }

        @Override
        public void run() {
            try {
                while (!stop.get()) {
                    sqsClient.sendMessage(new SendMessageRequest(queueUrl, theMessage));
                    producedCount.incrementAndGet();
                }
            } catch (AmazonClientException e) {
                // by default AmazonSQSClient retries calls 3 times before failing,
                // so, when this rare condition occurs, simply stop
                log.error("Producer: " + e.getMessage());
                System.exit(1);
            }
        }
    }

    /**
     * Producer thread using {@code SendMessageBatch} to send messages until
     * stopped.
     */
    private static class BatchProducer extends Thread {
        final AmazonSQS sqsClient;
        final String queueUrl;
        final int batchSize;
        final AtomicInteger producedCount;
        final AtomicBoolean stop;
        final String theMessage;

        BatchProducer(AmazonSQS sqsQueueBuffer, String queueUrl, int batchSize, int messageSizeByte,
                AtomicInteger producedCount, AtomicBoolean stop) {
            this.sqsClient = sqsQueueBuffer;
            this.queueUrl = queueUrl;
            this.batchSize = batchSize;
            this.producedCount = producedCount;
            this.stop = stop;
            this.theMessage = makeRandomString(messageSizeByte);
        }

        @Override
        public void run() {
            try {
                while (!stop.get()) {
                    SendMessageBatchRequest batchRequest = new SendMessageBatchRequest().withQueueUrl(queueUrl);

                    List<SendMessageBatchRequestEntry> entries = new ArrayList<SendMessageBatchRequestEntry>();
                    for (int i = 0; i < batchSize; i++)
                        entries.add(new SendMessageBatchRequestEntry().withId(Integer.toString(i))
                                .withMessageBody(theMessage));
                    batchRequest.setEntries(entries);

                    SendMessageBatchResult batchResult = sqsClient.sendMessageBatch(batchRequest);
                    producedCount.addAndGet(batchResult.getSuccessful().size());

                    // sendMessageBatch can return successfully, and yet individual batch
                    // items fail. So, make sure to retry the failed ones.
                    if (!batchResult.getFailed().isEmpty()) {
                        log.warn("Producer: retrying sending " + batchResult.getFailed().size() + " messages");
                        for (int i = 0, n = batchResult.getFailed().size(); i < n; i++) {
                            sqsClient.sendMessage(new SendMessageRequest(queueUrl, theMessage));
                            producedCount.incrementAndGet();
                        }
                    }
                }
            } catch (AmazonClientException e) {
                // by default AmazonSQSClient retries calls 3 times before failing,
                // so, when this rare condition occurs, simply stop
                log.error("BatchProducer: " + e.getMessage());
                System.exit(1);
            }
        }
    }

    /**
     * Consumer thread using {@code ReceiveMessage} and {@code DeleteMessage} to
     * consume messages until stopped.
     */
    private static class Consumer extends Thread {
        final AmazonSQS sqsClient;
        final String queueUrl;
        final AtomicInteger consumedCount;
        final AtomicBoolean stop;

        Consumer(AmazonSQS sqsClient, String queueUrl, AtomicInteger consumedCount, AtomicBoolean stop) {
            this.sqsClient = sqsClient;
            this.queueUrl = queueUrl;
            this.consumedCount = consumedCount;
            this.stop = stop;
        }

        @Override
        public void run() {
            try {
                ReceiveMessageResult result = null;
                Message m;
                while (!stop.get()) {
                    try {
                        result = sqsClient.receiveMessage(new ReceiveMessageRequest(queueUrl));

                        if (!result.getMessages().isEmpty()) {
                            m = result.getMessages().get(0);
                            sqsClient.deleteMessage(new DeleteMessageRequest(queueUrl, m.getReceiptHandle()));
                            consumedCount.incrementAndGet();
                        }
                    } catch (AmazonClientException e) {
                        log.error(e.getMessage());
                    }
                }
            } catch (AmazonClientException e) {
                // by default AmazonSQSClient retries calls 3 times before failing,
                // so, when this rare condition occurs, simply stop
                log.error("Consumer: " + e.getMessage());
                System.exit(1);
            }
        }
    }

    /**
     * Consumer thread using {@code ReceiveMessage} and {@code DeleteMessageBatch}
     * to consume messages until stopped.
     */
    private static class BatchConsumer extends Thread {
        final AmazonSQS sqsClient;
        final String queueUrl;
        final int batchSize;
        final AtomicInteger consumedCount;
        final AtomicBoolean stop;

        BatchConsumer(AmazonSQS sqsClient, String queueUrl, int batchSize, AtomicInteger consumedCount,
                AtomicBoolean stop) {
            this.sqsClient = sqsClient;
            this.queueUrl = queueUrl;
            this.batchSize = batchSize;
            this.consumedCount = consumedCount;
            this.stop = stop;
        }

        @Override
        public void run() {
            try {
                ReceiveMessageResult result = null;
                List<Message> messages;
                while (!stop.get()) {
                    result = sqsClient
                            .receiveMessage(new ReceiveMessageRequest(queueUrl).withMaxNumberOfMessages(batchSize));

                    if (!result.getMessages().isEmpty()) {
                        messages = result.getMessages();
                        DeleteMessageBatchRequest batchRequest = new DeleteMessageBatchRequest()
                                .withQueueUrl(queueUrl);

                        List<DeleteMessageBatchRequestEntry> entries = new ArrayList<DeleteMessageBatchRequestEntry>();
                        for (int i = 0, n = messages.size(); i < n; i++)
                            entries.add(new DeleteMessageBatchRequestEntry().withId(Integer.toString(i))
                                    .withReceiptHandle(messages.get(i).getReceiptHandle()));
                        batchRequest.setEntries(entries);

                        DeleteMessageBatchResult batchResult = sqsClient.deleteMessageBatch(batchRequest);
                        consumedCount.addAndGet(batchResult.getSuccessful().size());

                        // deleteMessageBatch can return successfully, and yet individual
                        // batch items fail. So, make sure to retry the failed ones.
                        if (!batchResult.getFailed().isEmpty()) {
                            int n = batchResult.getFailed().size();
                            log.warn("Producer: retrying deleting " + n + " messages");
                            for (BatchResultErrorEntry e : batchResult.getFailed()) {
                                sqsClient.deleteMessage(new DeleteMessageRequest(queueUrl,
                                        messages.get(Integer.parseInt(e.getId())).getReceiptHandle()));
                                consumedCount.incrementAndGet();
                            }
                        }
                    }
                }
            } catch (AmazonClientException e) {
                // by default AmazonSQSClient retries calls 3 times before failing,
                // so, when this rare condition occurs, simply stop
                log.error("BatchConsumer: " + e.getMessage());
                System.exit(1);
            }
        }
    }

    /**
     * Thread that prints, every second, the number of messages produced and
     * consumed so far.
     */
    private static class Monitor extends Thread {
        private final AtomicInteger producedCount;
        private final AtomicInteger consumedCount;
        private final AtomicBoolean stop;

        Monitor(AtomicInteger producedCount, AtomicInteger consumedCount, AtomicBoolean stop) {
            this.producedCount = producedCount;
            this.consumedCount = consumedCount;
            this.stop = stop;
        }

        @Override
        public void run() {
            try {
                while (!stop.get()) {
                    Thread.sleep(1000);
                    log.info("produced messages = " + producedCount.get() + ", consumed messages = "
                            + consumedCount.get());
                }
            } catch (InterruptedException e) {
                // allow thread to exit
            }
        }
    }

    private static String makeRandomString(int sizeByte) {
        byte[] bs = new byte[(int) Math.ceil(sizeByte * 5 / 8)];
        new Random().nextBytes(bs);
        bs[0] = (byte) ((bs[0] | 64) & 127);
        return new BigInteger(bs).toString(32);
    }
}