001/**
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017package org.apache.activemq.transport.nio;
018
019import java.io.DataInputStream;
020import java.io.DataOutputStream;
021import java.io.EOFException;
022import java.io.IOException;
023import java.net.Socket;
024import java.net.URI;
025import java.net.UnknownHostException;
026import java.nio.ByteBuffer;
027import java.nio.channels.SelectionKey;
028import java.nio.channels.SocketChannel;
029
030import javax.net.SocketFactory;
031
032import org.apache.activemq.MaxFrameSizeExceededException;
033import org.apache.activemq.openwire.OpenWireFormat;
034import org.apache.activemq.transport.Transport;
035import org.apache.activemq.transport.tcp.TcpTransport;
036import org.apache.activemq.util.IOExceptionSupport;
037import org.apache.activemq.util.ServiceStopper;
038import org.apache.activemq.wireformat.WireFormat;
039
040/**
041 * An implementation of the {@link Transport} interface using raw tcp/ip
042 *
043 *
044 */
045public class NIOTransport extends TcpTransport {
046
047    // private static final Logger log = LoggerFactory.getLogger(NIOTransport.class);
048    protected SocketChannel channel;
049    protected SelectorSelection selection;
050    protected ByteBuffer inputBuffer;
051    protected ByteBuffer currentBuffer;
052    protected int nextFrameSize;
053
054    public NIOTransport(WireFormat wireFormat, SocketFactory socketFactory, URI remoteLocation, URI localLocation) throws UnknownHostException, IOException {
055        super(wireFormat, socketFactory, remoteLocation, localLocation);
056    }
057
058    public NIOTransport(WireFormat wireFormat, Socket socket) throws IOException {
059        super(wireFormat, socket);
060    }
061
062    /**
063     * @param format
064     * @param socket
065     * @param initBuffer
066     * @throws IOException
067     */
068    public NIOTransport(WireFormat format, Socket socket, InitBuffer initBuffer) throws IOException {
069        super(format, socket, initBuffer);
070    }
071
072    @Override
073    protected void initializeStreams() throws IOException {
074        channel = socket.getChannel();
075        channel.configureBlocking(false);
076
077        // listen for events telling us when the socket is readable.
078        selection = SelectorManager.getInstance().register(channel, new SelectorManager.Listener() {
079            @Override
080            public void onSelect(SelectorSelection selection) {
081                serviceRead();
082            }
083
084            @Override
085            public void onError(SelectorSelection selection, Throwable error) {
086                if (error instanceof IOException) {
087                    onException((IOException)error);
088                } else {
089                    onException(IOExceptionSupport.create(error));
090                }
091            }
092        });
093
094        // Send the data via the channel
095        // inputBuffer = ByteBuffer.allocateDirect(8*1024);
096        inputBuffer = ByteBuffer.allocateDirect(getIoBufferSize());
097        currentBuffer = inputBuffer;
098        nextFrameSize = -1;
099        currentBuffer.limit(4);
100        NIOOutputStream outPutStream = new NIOOutputStream(channel, getIoBufferSize());
101        this.dataOut = new DataOutputStream(outPutStream);
102        this.buffOut = outPutStream;
103    }
104
105    protected int readFromBuffer() throws IOException {
106        return channel.read(currentBuffer);
107    }
108
109    protected void serviceRead() {
110        try {
111            while (true) {
112                //If the transport was already stopped then break
113                if (this.isStopped()) {
114                    return;
115                }
116
117                int readSize = readFromBuffer();
118                if (readSize == -1) {
119                    onException(new EOFException());
120                    selection.close();
121                    break;
122                }
123                if (readSize == 0) {
124                    break;
125                }
126
127                this.receiveCounter += readSize;
128                if (currentBuffer.hasRemaining()) {
129                    continue;
130                }
131
132                // Are we trying to figure out the size of the next frame?
133                if (nextFrameSize == -1) {
134                    assert inputBuffer == currentBuffer;
135
136                    // If the frame is too big to fit in our direct byte buffer,
137                    // Then allocate a non direct byte buffer of the right size
138                    // for it.
139                    inputBuffer.flip();
140                    nextFrameSize = inputBuffer.getInt() + 4;
141
142                    if (wireFormat instanceof OpenWireFormat) {
143                        OpenWireFormat openWireFormat = (OpenWireFormat)wireFormat;
144                        long maxFrameSize = openWireFormat.getMaxFrameSize();
145
146                        if (openWireFormat.isMaxFrameSizeEnabled() && nextFrameSize > maxFrameSize) {
147                            throw new MaxFrameSizeExceededException("Frame size of " + (nextFrameSize / (1024 * 1024)) + " MB larger than max allowed " + (maxFrameSize / (1024 * 1024)) + " MB");
148                        }
149                    }
150
151                    if (nextFrameSize > inputBuffer.capacity()) {
152                        currentBuffer = ByteBuffer.allocateDirect(nextFrameSize);
153                        currentBuffer.putInt(nextFrameSize);
154                    } else {
155                        inputBuffer.limit(nextFrameSize);
156                    }
157
158                } else {
159                    currentBuffer.flip();
160
161                    Object command = wireFormat.unmarshal(new DataInputStream(new NIOInputStream(currentBuffer)));
162                    doConsume(command);
163
164                    nextFrameSize = -1;
165                    inputBuffer.clear();
166                    inputBuffer.limit(4);
167                    currentBuffer = inputBuffer;
168                }
169
170            }
171
172        } catch (IOException e) {
173            onException(e);
174        } catch (Throwable e) {
175            onException(IOExceptionSupport.create(e));
176        }
177    }
178
179    @Override
180    protected void doStart() throws Exception {
181        connect();
182        selection.setInterestOps(SelectionKey.OP_READ);
183        selection.enable();
184    }
185
186    @Override
187    protected void doStop(ServiceStopper stopper) throws Exception {
188        if (selection != null) {
189            selection.close();
190            selection = null;
191        }
192        super.doStop(stopper);
193    }
194}