SQSMessageRunnerService

package com.atlassian.messagequeue.internal.sqs;

import com.amazonaws.AbortedException;
import com.amazonaws.http.timers.client.SdkInterruptedException;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.amazonaws.services.sqs.model.InvalidMessageContentsException;
import com.amazonaws.services.sqs.model.ReceiveMessageRequest;
import com.amazonaws.services.sqs.model.ReceiveMessageResult;
import com.amazonaws.services.sqs.model.SendMessageRequest;
import com.amazonaws.services.sqs.model.SendMessageResult;
import com.amazonaws.services.sqs.model.UnsupportedOperationException;
import com.atlassian.messagequeue.Message;
import com.atlassian.messagequeue.MessageInformationService;
import com.atlassian.messagequeue.MessageRunnerNotRegisteredException;
import com.atlassian.messagequeue.MessageRunnerService;
import com.atlassian.messagequeue.MessageRunnerServiceException;
import com.atlassian.messagequeue.MessageSerializationException;
import com.atlassian.messagequeue.internal.core.NestedMessage;
import com.atlassian.messagequeue.internal.core.NestedMessageConsumer;
import com.atlassian.messagequeue.internal.core.NestedMessageSerializer;
import com.atlassian.messagequeue.internal.core.MessageRunnerRegistryHelper;
import com.atlassian.tenant.impl.TenantIdSetter;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import org.slf4j.MDC;

import javax.annotation.Nonnull;
import javax.annotation.PostConstruct;
import javax.annotation.PreDestroy;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.Future;
import java.util.concurrent.LinkedBlockingQueue;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;

import static com.atlassian.messagequeue.internal.core.NestedMessageConsumer.MDC_MESSAGE_RUNNER_KEY;
import static java.util.Objects.requireNonNull;

/**
 * SQS message runner service that produces messages onto an SQS queue and manages a pool of threads that consume from said queue.
 * <p>
 * To properly initialise an instance of this service, a client must call {@link #initialiseMessageConsumers()} after
 * constructing an instance (or else ensure the dependency injection framework will automatically invoke methods
 * annotated with @PostConstruct).
 * <p>
 * To properly dispose of an instance of this service, a client must call {@link #shutdown()} (or else ensure the
 * dependency injection framework will automatically invoke methods annotated with @PreDestroy).
 */
public class SQSMessageRunnerService implements MessageRunnerService {
    static final int MAX_CONCURRENT_CONSUMERS = 2 * Runtime.getRuntime().availableProcessors();
    static final String SENT_TIMESTAMP = "SentTimestamp";
    static final long AWAIT_TERMINATION_TIMEOUT_SECONDS = Long.getLong("amq.sqs.await.termination.timeout", TimeUnit.SECONDS.convert(1, TimeUnit.HOURS));

    private static final Logger log = LoggerFactory.getLogger(SQSMessageRunnerService.class);
    private static final String ALL = "ALL";

    private static final int CONCURRENT_CONSUMERS = Integer.getInteger("amq.sqs.concurrent.consumers", Runtime.getRuntime().availableProcessors());
    private static final int DEFAULT_RECEIVE_WAIT_TIME_SECONDS = Integer.getInteger("amq.sqs.receive.wait.time", 20);
    private static final String MDC_MESSAGE_ID = "AMQ-messageID";
    private static final String MDC_MESSAGE_SENT_TIMESTAMP = "AMQ-messageSentTimestamp";

    private final AmazonSQSClient amazonSQSClient;
    private final MessageRunnerRegistryHelper messageRunnerRegistryHelper;
    private final AtomicBoolean shuttingDown = new AtomicBoolean(false);
    private final int receiveWaitTimeSeconds;
    private final String queueUrl;
    private final int concurrentConsumers;
    private final ExecutorService threadPool;
    private final MessageInformationService messageInformationService;
    private final NestedMessageSerializer nestedMessageSerializer;
    private final NestedMessageConsumer nestedMessageConsumer;
    private final SQSMessageVisibilityTimeoutManager sqsMessageVisibilityTimeoutManager;

    public SQSMessageRunnerService(AmazonSQSClient amazonSQSClient, String queueUrl,
                                   MessageRunnerRegistryHelper messageRunnerRegistryHelper,
                                   TenantIdSetter tenantIdSetter,
                                   MessageInformationService messageInformationService, NestedMessageSerializer nestedMessageSerializer) {
        this(amazonSQSClient, queueUrl, DEFAULT_RECEIVE_WAIT_TIME_SECONDS,
                CONCURRENT_CONSUMERS, messageRunnerRegistryHelper, tenantIdSetter, messageInformationService, nestedMessageSerializer,
                new SQSMessageVisibilityTimeoutManager(amazonSQSClient, queueUrl));
    }

    /**
     * Package private visibility for testing.
     */
    SQSMessageRunnerService(AmazonSQSClient amazonSQSClient, String queueUrl, int receiveWaitTimeSeconds,
                            int concurrentConsumers,
                            MessageRunnerRegistryHelper messageRunnerRegistryHelper,
                            TenantIdSetter tenantIdSetter,
                            MessageInformationService messageInformationService,
                            NestedMessageSerializer nestedMessageSerializer,
                            SQSMessageVisibilityTimeoutManager sqsMessageVisibilityTimeoutManager) {
        this.amazonSQSClient = requireNonNull(amazonSQSClient);
        this.messageRunnerRegistryHelper = requireNonNull(messageRunnerRegistryHelper);
        this.receiveWaitTimeSeconds = receiveWaitTimeSeconds;
        this.messageInformationService = requireNonNull(messageInformationService);
        this.nestedMessageSerializer = requireNonNull(nestedMessageSerializer);
        this.nestedMessageConsumer = new NestedMessageConsumer(messageRunnerRegistryHelper, tenantIdSetter);

        if (concurrentConsumers <= 0 || concurrentConsumers > MAX_CONCURRENT_CONSUMERS) {
            throw new IllegalArgumentException("concurrent consumers must be > 0 and <= " + MAX_CONCURRENT_CONSUMERS + ". Received: " + concurrentConsumers);
        }
        this.concurrentConsumers = concurrentConsumers;
        this.queueUrl = requireNonNull(queueUrl, "queueUrl");

        final DefaultThreadFactory threadFactory = new DefaultThreadFactory("sqs-consumer-thread-%d", (t, throwable) -> {
            log.warn("SQS consumer thread '{}' died due to an exception.", t.getName(), throwable);
        });

        this.threadPool = new ThreadPoolExecutor(concurrentConsumers, concurrentConsumers, 0L, TimeUnit.MILLISECONDS,
                new LinkedBlockingQueue<>(),
                threadFactory) {
            @Override
            protected void afterExecute(Runnable r, Throwable t) {
                if (t != null && !this.isShutdown()) {
                    this.execute(r); // we want another thread to resume executing this runnable
                }
            }
        };

        this.sqsMessageVisibilityTimeoutManager = requireNonNull(sqsMessageVisibilityTimeoutManager, "sqsMessageVisibilityTimeoutManager");

        log.info("Constructing {} (concurrentConsumers: {}, queueUrl: {}, receiveWaitTimeSeconds: {})",
                this.getClass().getSimpleName(), concurrentConsumers, queueUrl, receiveWaitTimeSeconds);
    }

    @PostConstruct
    public void initialiseMessageConsumers() {
        log.info("Initialising SQS messsage consumers");
        for (int i = 0; i < concurrentConsumers; i++) {
            threadPool.execute(new SQSMessageConsumer());
        }
    }

    private class SQSMessageConsumer implements Runnable {
        @Override
        public void run() {
            while (!Thread.currentThread().isInterrupted() && !shuttingDown.get()) {
                final ReceiveMessageRequest receiveMessageRequest = new ReceiveMessageRequest(queueUrl)
                        .withMaxNumberOfMessages(1)
                        .withAttributeNames(SENT_TIMESTAMP)
                        .withMessageAttributeNames(ALL)
                        .withWaitTimeSeconds(receiveWaitTimeSeconds)
                        .withVisibilityTimeout(sqsMessageVisibilityTimeoutManager.getVisibilityTimeoutSeconds());

                final ReceiveMessageResult receiveMessageResult;
                try {
                    receiveMessageResult = amazonSQSClient.receiveMessage(receiveMessageRequest);
                    processReceiveMessageResult(receiveMessageResult);
                } catch (Throwable t) {
                    if (isThrowbableThrownForInterrupt(t)) {
                        break;
                    } else {
                        log.error("Error occurred while consuming a message from SQS", t);
                    }
                }
            }
        }

        private void processReceiveMessageResult(ReceiveMessageResult receiveMessageResult) {
            if (receiveMessageResult.getMessages().isEmpty()) {
                return;
            }

            if (receiveMessageResult.getMessages().size() > 1) {
                throw new AssertionError("Number of messages received greater than the max number specified in the request.");
            }

            final com.amazonaws.services.sqs.model.Message message = receiveMessageResult.getMessages().get(0);

            final Future<?> visibilityTimeoutExtensionFuture = sqsMessageVisibilityTimeoutManager.scheduleVisibilityTimeoutExtension(message.getReceiptHandle());
            try (MDC.MDCCloseable mdcMessageId = MDC.putCloseable(MDC_MESSAGE_ID, message.getMessageId())) {
                if (log.isInfoEnabled()) {
                    log.info("Consuming message from {} (messageId: {}, receiptHandle: {})", queueUrl, message.getMessageId(), message.getReceiptHandle());
                }

                final NestedMessage nestedMessage;
                try {
                    nestedMessage = nestedMessageSerializer.deserialize(message.getBody());
                } catch (MessageSerializationException e) {
                    log.error("Message received could not be deserialized: {}. Message will be deleted (messageID: {})", e.getMessage(), message.getMessageId());
                    amazonSQSClient.deleteMessage(queueUrl, message.getReceiptHandle());
                    return;
                }

                try (MDC.MDCCloseable mdcMessageSentTimestamp = MDC.putCloseable(MDC_MESSAGE_SENT_TIMESTAMP, message.getAttributes().get(SENT_TIMESTAMP))) {
                    nestedMessageConsumer.consume(nestedMessage, new SQSMessageContext(message.getMessageId(),
                            nestedMessage.getPayload(), amazonSQSClient, queueUrl, message.getReceiptHandle(), shuttingDown, visibilityTimeoutExtensionFuture));
                }
            } finally {
                if (!visibilityTimeoutExtensionFuture.isDone()) {
                    final boolean cancelled = visibilityTimeoutExtensionFuture.cancel(true);
                    log.info("Cancelled extension of visibility timeout. Cancellation status: {}", cancelled);
                }
            }
        }
    }

    private static boolean isThrowbableThrownForInterrupt(Throwable throwable) {
        return throwable.getCause() instanceof SdkInterruptedException
                || throwable instanceof AbortedException; // thrown by aws SDK InputStream classes in response to an interrupt
    }

    @PreDestroy
    public void shutdown() {
        log.info("Shutting down {}", this.getClass().getSimpleName());

        shuttingDown.compareAndSet(false, true);

        threadPool.shutdown();
        try {
            if (!threadPool.awaitTermination(AWAIT_TERMINATION_TIMEOUT_SECONDS, TimeUnit.SECONDS)) {
                log.warn("Pool did not terminate in {} seconds", AWAIT_TERMINATION_TIMEOUT_SECONDS);
            }
        } catch (InterruptedException ie) {
            Thread.currentThread().interrupt();
        }

        sqsMessageVisibilityTimeoutManager.shutdown();
    }

    @Override
    public void addMessage(@Nonnull Message message) {
        if (!messageRunnerRegistryHelper.getMessageRunner(message.getRunnerKey()).isPresent()) {
            throw new MessageRunnerNotRegisteredException(message.getRunnerKey());
        }

        final SendMessageResult sendMessageResult;
        try {
            final String queueUrl = messageInformationService.getQueueUrl(message.getRunnerKey());

            sendMessageResult = amazonSQSClient.sendMessage(new SendMessageRequest(queueUrl,
                    messageInformationService.toPayload(message)));
        } catch (InvalidMessageContentsException e) {
            throw new MessageRunnerServiceException(e);
        } catch (UnsupportedOperationException e) {
            throw new java.lang.UnsupportedOperationException(e);
        }

        if (log.isInfoEnabled()) {
            try (MDC.MDCCloseable mdcCloseable1 = MDC.putCloseable(MDC_MESSAGE_ID, sendMessageResult.getMessageId());
                 MDC.MDCCloseable mdcCloseable2 = MDC.putCloseable(MDC_MESSAGE_RUNNER_KEY, message.getRunnerKey().toString())) {
                log.info("Message produced to {} (messageId: {}, messageRunnerKey: {})",
                        queueUrl, sendMessageResult.getMessageId(), message.getRunnerKey().toString());
            }
        }
    }
}