SQSNotice.java

/*
 * Copyright 2013 Netflix, Inc.
 *
 *    Licensed 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 com.netflix.suro.sink.notice;

import com.amazonaws.ClientConfiguration;
import com.amazonaws.auth.AWSCredentialsProvider;
import com.amazonaws.services.sqs.AmazonSQSClient;
import com.amazonaws.services.sqs.model.*;
import com.fasterxml.jackson.annotation.JacksonInject;
import com.fasterxml.jackson.annotation.JsonCreator;
import com.fasterxml.jackson.annotation.JsonProperty;
import com.google.common.base.Charsets;
import com.google.common.base.Joiner;
import com.google.common.base.Preconditions;
import com.netflix.servo.annotations.DataSourceType;
import com.netflix.servo.annotations.Monitor;
import com.netflix.servo.monitor.Monitors;
import com.netflix.suro.TagKey;
import com.netflix.util.Pair;
import org.apache.commons.codec.binary.Base64;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.atomic.AtomicLong;

/**
 * SQS {@link com.netflix.suro.sink.notice.Notice} implementation
 *
 * @author jbae
 */
public class SQSNotice implements Notice<String> {
    static Logger log = LoggerFactory.getLogger(SQSNotice.class);

    public static final String TYPE = "sqs";

    private final List<String> queues;
    private final List<String> queueUrls = new ArrayList<String>();
    private final boolean enableBase64Encoding;

    private AmazonSQSClient sqsClient;
    private final AWSCredentialsProvider credentialsProvider;

    private ClientConfiguration clientConfig;
    private final String region;

    @Monitor(name = TagKey.SENT_COUNT, type = DataSourceType.COUNTER)
    private AtomicLong sentMessageCount = new AtomicLong(0);
    @Monitor(name = TagKey.LOST_COUNT, type = DataSourceType.COUNTER)
    private AtomicLong lostMessageCount = new AtomicLong(0);
    @Monitor(name = TagKey.RECV_COUNT, type = DataSourceType.COUNTER)
    private AtomicLong recvMessageCount = new AtomicLong(0);

    @JsonCreator
    public SQSNotice(
            @JsonProperty("queues") List<String> queues,
            @JsonProperty("region") @JacksonInject("region") String region,
            @JsonProperty("connectionTimeout") int connectionTimeout,
            @JsonProperty("maxConnections") int maxConnections,
            @JsonProperty("socketTimeout") int socketTimeout,
            @JsonProperty("maxRetries") int maxRetries,
            @JsonProperty("enableBase64Encoding") boolean enableBase64Encoding,
            @JacksonInject AmazonSQSClient sqsClient,
            @JacksonInject AWSCredentialsProvider credentialsProvider) {
        this.queues = queues;
        this.region = region;

        this.enableBase64Encoding = enableBase64Encoding;
        this.sqsClient = sqsClient;
        this.credentialsProvider = credentialsProvider;

        Preconditions.checkArgument(queues.size() > 0);
        Preconditions.checkNotNull(region);

        clientConfig = new ClientConfiguration();
        if (connectionTimeout > 0) {
            clientConfig = clientConfig.withConnectionTimeout(connectionTimeout);
        }
        if (maxConnections > 0) {
            clientConfig = clientConfig.withMaxConnections(maxConnections);
        }
        if (socketTimeout > 0) {
            clientConfig = clientConfig.withSocketTimeout(socketTimeout);
        }
        if (maxRetries > 0) {
            clientConfig = clientConfig.withMaxErrorRetry(maxRetries);
        }

        Monitors.registerObject(Joiner.on('_').join(queues), this);
    }

    @Override
    public void init() {
        if (sqsClient == null) { // not injected
            sqsClient = new AmazonSQSClient(credentialsProvider, clientConfig);
        }
        String endpoint = "sqs." + this.region + ".amazonaws.com";
        sqsClient.setEndpoint(endpoint);

        for (String queueName : queues) {
            GetQueueUrlRequest request = new GetQueueUrlRequest();
            request.setQueueName(queueName);
            queueUrls.add(sqsClient.getQueueUrl(request).getQueueUrl());
        }

        log.info(String.format("SQSNotice initialized with the endpoint: %s, queue: %s",
                endpoint, queues));
    }

    @Override
    public boolean send(String message) {
        boolean sent = false;

        try {
            for (String queueUrl : queueUrls) {
                SendMessageRequest request = new SendMessageRequest()
                        .withQueueUrl(queueUrl);

                if(enableBase64Encoding) {
                    request = request.withMessageBody(
                            new String(
                                    Base64.encodeBase64(
                                            message.getBytes(Charsets.UTF_8))));
                } else {
                    request = request.withMessageBody(message);
                }
                sqsClient.sendMessage(request);
                log.info("SQSNotice: " + message + " sent to " + queueUrl);
                if (!sent) {
                    sentMessageCount.incrementAndGet();
                    sent = true;
                }
            }
        } catch (Exception e) {
            log.error("Exception while sending SQS notice: " + e.getMessage(), e);
        }

        if (!sent) {
            lostMessageCount.incrementAndGet();
        }

        return sent;
    }

    @Override
    public String recv() {
        ReceiveMessageRequest request = new ReceiveMessageRequest()
                .withQueueUrl(queueUrls.get(0))
                .withMaxNumberOfMessages(1);

        try {
            ReceiveMessageResult result = sqsClient.receiveMessage(request);
            if (!result.getMessages().isEmpty()) {
                Message msg = result.getMessages().get(0);

                recvMessageCount.incrementAndGet();

                DeleteMessageRequest deleteReq = new DeleteMessageRequest()
                        .withQueueUrl(queueUrls.get(0))
                        .withReceiptHandle(msg.getReceiptHandle());
                sqsClient.deleteMessage(deleteReq);

                if (enableBase64Encoding) {
                    return new String(
                            Base64.decodeBase64(msg.getBody().getBytes()),
                            Charsets.UTF_8);
                } else {
                    return msg.getBody();
                }
            } else {
                return "";
            }
        } catch (Exception e) {
            log.error("Exception while recving SQS notice: " + e.getMessage(), e);
            return "";
        }
    }

    @Override
    public Pair<String, String> peek() {
        ReceiveMessageRequest request = new ReceiveMessageRequest()
                .withQueueUrl(queueUrls.get(0))
                .withMaxNumberOfMessages(1);

        try {
            ReceiveMessageResult result = sqsClient.receiveMessage(request);
            if (!result.getMessages().isEmpty()) {
                Message msg = result.getMessages().get(0);

                recvMessageCount.incrementAndGet();

                if (enableBase64Encoding) {
                    return new Pair<String, String>(
                            msg.getReceiptHandle(),
                            new String(
                                    Base64.decodeBase64(msg.getBody().getBytes()),
                                    Charsets.UTF_8));
                } else {
                    return new Pair<String, String>(
                            msg.getReceiptHandle(),
                            msg.getBody());
                }
            } else {
                return null;
            }
        } catch (Exception e) {
            log.error("Exception while recving SQS notice: " + e.getMessage(), e);
            return null;
        }
    }

    @Override
    public void remove(String key) {
        DeleteMessageRequest deleteReq = new DeleteMessageRequest()
                .withQueueUrl(queueUrls.get(0))
                .withReceiptHandle(key);
        sqsClient.deleteMessage(deleteReq);
    }

    @Override
    public String getStat() {
        return String.format("SQSNotice with the queues: %s, sent : %d, received: %d, dropped: %d",
                queues, sentMessageCount.get(), recvMessageCount.get(), lostMessageCount.get());
    }
}