MessageSetBuilder.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.message;

import com.google.common.io.ByteArrayDataOutput;
import com.netflix.suro.ClientConfig;
import com.netflix.suro.client.async.Queue4Client;
import com.netflix.suro.thrift.TMessageSet;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.List;
import java.util.zip.CRC32;

/**
 * The payload for Suro's Thrift communication is {@link TMessageSet}, a thrift presentation of
 * a set of messages. This class can be helpful to easily create {@link TMessageSet} instances.
 *
 * @author jbae
 */
public class MessageSetBuilder {
    private static final Logger log = LoggerFactory.getLogger(MessageSetBuilder.class);

    private final ClientConfig config;
    private List<Message> messageList;
    private Compression compression = Compression.LZF;

    /**
     * @param config contains information including application name, etc
     */
    public MessageSetBuilder(ClientConfig config) {
        this.config = config;
        messageList = new ArrayList<Message>();
    }

    public MessageSetBuilder withMessage(String routingKey, byte[] payload) {
        this.messageList.add(new Message(routingKey, payload));
        return this;
    }

    public MessageSetBuilder withCompression(Compression compresson) {
        this.compression = compresson;
        return this;
    }

    public TMessageSet build() {
        try {
            byte[] buffer = createPayload(messageList, compression);
            long crc = getCRC(buffer);


            return new TMessageSet(
                    config.getApp(),
                    messageList.size(),
                    (byte) compression.getId(),
                    crc,
                    ByteBuffer.wrap(buffer));
        } catch (IOException e) {
            log.error("Exception on building TMessageSet: " + e.getMessage(), e);
            return null;
        } finally {
            messageList.clear();
        }
    }

    /**
     * @return number of messages in MessageSet
     */
    public int size() {
        return messageList.size();
    }

    /**
     * Create compressed byte[] from the list of messages. Each message contains
     * byte[] as its message body, so, this method is simply flattening byte[]
     * for all messages in the messageList
     *
     * @param messageList a list of messages for payload
     * @param compression Compression method to be applied to the payload
     * @return A byte array that encodes the build payload
     * @throws IOException
     */
    public static byte[] createPayload(List<Message> messageList, Compression compression) throws IOException {
        ByteArrayDataOutput out = new ByteArrayDataOutputStream(outputStream.get());
        for (Message message : messageList) {
            message.write(out);
        }

        return compression.compress(out.toByteArray());
    }

    private static ThreadLocal<ByteArrayOutputStream> outputStream =
            new ThreadLocal<ByteArrayOutputStream>() {
                @Override
                protected ByteArrayOutputStream initialValue() {
                    return new ByteArrayOutputStream();
                }

                @Override
                public ByteArrayOutputStream get() {
                    ByteArrayOutputStream b = super.get();
                    b.reset();
                    return b;
                }
            };

    /**
     * Compute CRC32 value for byte[]
     *
     * @param buffer all the bytes in the buffer will be used for CRC32 calculation
     *
     * @return a CRC32 value for the given byte array
     */
    public static long getCRC(byte[] buffer) {
        CRC32 crc = new CRC32();
        crc.update(buffer);
        return crc.getValue();
    }

    /**
     * Drains the given number of messages from the givne queue. Instead of calling {@link #withMessage(String, byte[])},
     * we can call this method.
     * This method is reverse one of JDK BlockingQueue.drainTo.
     *
     * @param queue the queue to drain messages from
     * @param size the number of messages to drain from the given queue
     */
    public void drainFrom(Queue4Client queue, int size) {
        queue.drain(size, messageList);
    }
}