trunk/activemq-core/src/main/java/org/apache/activemq/transport/udp/CommandDatagramChannel.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.activemq.transport.udp;

import java.io.DataInputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.DatagramChannel;

import org.apache.activemq.command.Command;
import org.apache.activemq.command.Endpoint;
import org.apache.activemq.command.LastPartialCommand;
import org.apache.activemq.command.PartialCommand;
import org.apache.activemq.openwire.BooleanStream;
import org.apache.activemq.openwire.OpenWireFormat;
import org.apache.activemq.transport.reliable.ReplayBuffer;
import org.apache.activemq.util.ByteArrayInputStream;
import org.apache.activemq.util.ByteArrayOutputStream;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * A strategy for reading datagrams and de-fragmenting them together.
 * 
 * 
 */
public class CommandDatagramChannel extends CommandChannelSupport {

    private static final Logger LOG = LoggerFactory.getLogger(CommandDatagramChannel.class);

    private DatagramChannel channel;
    private ByteBufferPool bufferPool;

    // reading
    private Object readLock = new Object();
    private ByteBuffer readBuffer;

    // writing
    private Object writeLock = new Object();
    private int defaultMarshalBufferSize = 64 * 1024;
    private volatile int receiveCounter;

    public CommandDatagramChannel(UdpTransport transport, OpenWireFormat wireFormat, int datagramSize, SocketAddress targetAddress, DatagramHeaderMarshaller headerMarshaller,
                                  DatagramChannel channel, ByteBufferPool bufferPool) {
        super(transport, wireFormat, datagramSize, targetAddress, headerMarshaller);
        this.channel = channel;
        this.bufferPool = bufferPool;
    }

    public void start() throws Exception {
        bufferPool.setDefaultSize(datagramSize);
        bufferPool.start();
        readBuffer = bufferPool.borrowBuffer();
    }

    public void stop() throws Exception {
        bufferPool.stop();
    }

    public Command read() throws IOException {
        Command answer = null;
        Endpoint from = null;
        synchronized (readLock) {
            while (true) {
                readBuffer.clear();
                SocketAddress address = channel.receive(readBuffer);

                readBuffer.flip();

                if (readBuffer.limit() == 0) {
                    continue;
                }
                
                receiveCounter++;
                from = headerMarshaller.createEndpoint(readBuffer, address);

                int remaining = readBuffer.remaining();
                byte[] data = new byte[remaining];
                readBuffer.get(data);

                // TODO could use a DataInput implementation that talks direct
                // to
                // the ByteBuffer to avoid object allocation and unnecessary
                // buffering?
                DataInputStream dataIn = new DataInputStream(new ByteArrayInputStream(data));
                answer = (Command)wireFormat.unmarshal(dataIn);
                break;
            }
        }
        if (answer != null) {
            answer.setFrom(from);

            if (LOG.isDebugEnabled()) {
                LOG.debug("Channel: " + name + " received from: " + from + " about to process: " + answer);
            }
        }
        return answer;
    }

    public void write(Command command, SocketAddress address) throws IOException {
        synchronized (writeLock) {

            ByteArrayOutputStream largeBuffer = new ByteArrayOutputStream(defaultMarshalBufferSize);
            wireFormat.marshal(command, new DataOutputStream(largeBuffer));
            byte[] data = largeBuffer.toByteArray();
            int size = data.length;

            ByteBuffer writeBuffer = bufferPool.borrowBuffer();
            writeBuffer.clear();
            headerMarshaller.writeHeader(command, writeBuffer);

            if (size > writeBuffer.remaining()) {
                // lets split the command up into chunks
                int offset = 0;
                boolean lastFragment = false;
                int length = data.length;
                for (int fragment = 0; !lastFragment; fragment++) {
                    // write the header
                    if (fragment > 0) {
                        writeBuffer = bufferPool.borrowBuffer();
                        writeBuffer.clear();
                        headerMarshaller.writeHeader(command, writeBuffer);
                    }

                    int chunkSize = writeBuffer.remaining();

                    // we need to remove the amount of overhead to write the
                    // partial command

                    // lets write the flags in there
                    BooleanStream bs = null;
                    if (wireFormat.isTightEncodingEnabled()) {
                        bs = new BooleanStream();
                        bs.writeBoolean(true); // the partial data byte[] is
                        // never null
                    }

                    // lets remove the header of the partial command
                    // which is the byte for the type and an int for the size of
                    // the byte[]

                    // data type + the command ID + size of the partial data
                    chunkSize -= 1 + 4 + 4;

                    // the boolean flags
                    if (bs != null) {
                        chunkSize -= bs.marshalledSize();
                    } else {
                        chunkSize -= 1;
                    }

                    if (!wireFormat.isSizePrefixDisabled()) {
                        // lets write the size of the command buffer
                        writeBuffer.putInt(chunkSize);
                        chunkSize -= 4;
                    }

                    lastFragment = offset + chunkSize >= length;
                    if (chunkSize + offset > length) {
                        chunkSize = length - offset;
                    }

                    if (lastFragment) {
                        writeBuffer.put(LastPartialCommand.DATA_STRUCTURE_TYPE);
                    } else {
                        writeBuffer.put(PartialCommand.DATA_STRUCTURE_TYPE);
                    }

                    if (bs != null) {
                        bs.marshal(writeBuffer);
                    }

                    int commandId = command.getCommandId();
                    if (fragment > 0) {
                        commandId = sequenceGenerator.getNextSequenceId();
                    }
                    writeBuffer.putInt(commandId);
                    if (bs == null) {
                        writeBuffer.put((byte)1);
                    }

                    // size of byte array
                    writeBuffer.putInt(chunkSize);

                    // now the data
                    writeBuffer.put(data, offset, chunkSize);

                    offset += chunkSize;
                    sendWriteBuffer(commandId, address, writeBuffer, false);
                }
            } else {
                writeBuffer.put(data);
                sendWriteBuffer(command.getCommandId(), address, writeBuffer, false);
            }
        }
    }

    // Properties
    // -------------------------------------------------------------------------

    public ByteBufferPool getBufferPool() {
        return bufferPool;
    }

    /**
     * Sets the implementation of the byte buffer pool to use
     */
    public void setBufferPool(ByteBufferPool bufferPool) {
        this.bufferPool = bufferPool;
    }

    // Implementation methods
    // -------------------------------------------------------------------------
    protected void sendWriteBuffer(int commandId, SocketAddress address, ByteBuffer writeBuffer, boolean redelivery) throws IOException {
        // lets put the datagram into the replay buffer first to prevent timing
        // issues
        ReplayBuffer bufferCache = getReplayBuffer();
        if (bufferCache != null && !redelivery) {
            bufferCache.addBuffer(commandId, writeBuffer);
        }

        writeBuffer.flip();

        if (LOG.isDebugEnabled()) {
            String text = redelivery ? "REDELIVERING" : "sending";
            LOG.debug("Channel: " + name + " " + text + " datagram: " + commandId + " to: " + address);
        }
        channel.send(writeBuffer, address);
    }

    public void sendBuffer(int commandId, Object buffer) throws IOException {
        if (buffer != null) {
            ByteBuffer writeBuffer = (ByteBuffer)buffer;
            sendWriteBuffer(commandId, getReplayAddress(), writeBuffer, true);
        } else {
            if (LOG.isWarnEnabled()) {
                LOG.warn("Request for buffer: " + commandId + " is no longer present");
            }
        }
    }

    public int getReceiveCounter() {
        return receiveCounter;
    }

}