001    /*
002     * Copyright (C) 2011-2012, FuseSource Corp.  All rights reserved.
003     *
004     *     http://fusesource.com
005     *
006     * Licensed under the Apache License, Version 2.0 (the "License");
007     * you may not use this file except in compliance with the License.
008     * You may obtain a copy of the License at
009     *
010     *     http://www.apache.org/licenses/LICENSE-2.0
011     *
012     * Unless required by applicable law or agreed to in writing, software
013     * distributed under the License is distributed on an "AS IS" BASIS,
014     * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
015     * See the License for the specific language governing permissions and
016     * limitations under the License.
017     */
018    package org.fusesource.hawtdispatch.transport;
019    
020    import org.fusesource.hawtbuf.Buffer;
021    import org.fusesource.hawtbuf.DataByteArrayOutputStream;
022    
023    import java.io.EOFException;
024    import java.io.IOException;
025    import java.net.ProtocolException;
026    import java.net.SocketException;
027    import java.nio.ByteBuffer;
028    import java.nio.channels.GatheringByteChannel;
029    import java.nio.channels.ReadableByteChannel;
030    import java.nio.channels.SocketChannel;
031    import java.nio.channels.WritableByteChannel;
032    import java.util.LinkedList;
033    
034    /**
035     * Provides an abstract base class to make implementing the ProtocolCodec interface
036     * easier.
037     *
038     * @author <a href="http://hiramchirino.com">Hiram Chirino</a>
039     */
040    public abstract class AbstractProtocolCodec implements ProtocolCodec {
041    
042        protected int writeBufferSize = 1024 * 64;
043        protected long writeCounter = 0L;
044        protected GatheringByteChannel writeChannel = null;
045        protected DataByteArrayOutputStream nextWriteBuffer = new DataByteArrayOutputStream(writeBufferSize);
046        protected long lastWriteIoSize = 0;
047    
048        protected LinkedList<ByteBuffer> writeBuffer = new LinkedList<ByteBuffer>();
049        private long writeBufferRemaining = 0;
050    
051    
052        public static interface Action {
053            Object apply() throws IOException;
054        }
055    
056        protected long readCounter = 0L;
057        protected int readBufferSize = 1024 * 64;
058        protected ReadableByteChannel readChannel = null;
059        protected ByteBuffer readBuffer = ByteBuffer.allocate(readBufferSize);
060        protected ByteBuffer directReadBuffer = null;
061    
062        protected int readEnd;
063        protected int readStart;
064        protected int lastReadIoSize;
065        protected Action nextDecodeAction;
066        protected boolean trim = true;
067    
068    
069        public void setWritableByteChannel(WritableByteChannel channel) throws SocketException {
070            this.writeChannel = (GatheringByteChannel) channel;
071            if (this.writeChannel instanceof SocketChannel) {
072                writeBufferSize = ((SocketChannel) this.writeChannel).socket().getSendBufferSize();
073            } else if (this.writeChannel instanceof SslTransport.SSLChannel) {
074                writeBufferSize = ((SslTransport.SSLChannel) this.writeChannel).socket().getSendBufferSize();
075            }
076        }
077    
078        public int getReadBufferSize() {
079            return readBufferSize;
080        }
081    
082        public int getWriteBufferSize() {
083            return writeBufferSize;
084        }
085    
086        public boolean full() {
087            return writeBufferRemaining >= writeBufferSize;
088        }
089    
090        public boolean isEmpty() {
091            return writeBufferRemaining == 0 && nextWriteBuffer.size() == 0;
092        }
093    
094        public long getWriteCounter() {
095            return writeCounter;
096        }
097    
098        public long getLastWriteSize() {
099            return lastWriteIoSize;
100        }
101    
102        abstract protected void encode(Object value) throws IOException;
103    
104        public ProtocolCodec.BufferState write(Object value) throws IOException {
105            if (full()) {
106                return ProtocolCodec.BufferState.FULL;
107            } else {
108                boolean wasEmpty = isEmpty();
109                encode(value);
110                if (nextWriteBuffer.size() >= (writeBufferSize >> 1)) {
111                    flushNextWriteBuffer();
112                }
113                if (wasEmpty) {
114                    return ProtocolCodec.BufferState.WAS_EMPTY;
115                } else {
116                    return ProtocolCodec.BufferState.NOT_EMPTY;
117                }
118            }
119        }
120    
121        protected void writeDirect(ByteBuffer value) throws IOException {
122            // is the direct buffer small enough to just fit into the nextWriteBuffer?
123            int nextnextPospos = nextWriteBuffer.position();
124            int valuevalueLengthlength = value.remaining();
125            int available = nextWriteBuffer.getData().length - nextnextPospos;
126            if (available > valuevalueLengthlength) {
127                value.get(nextWriteBuffer.getData(), nextnextPospos, valuevalueLengthlength);
128                nextWriteBuffer.position(nextnextPospos + valuevalueLengthlength);
129            } else {
130                if (nextWriteBuffer.size() != 0) {
131                    flushNextWriteBuffer();
132                }
133                writeBuffer.add(value);
134                writeBufferRemaining += value.remaining();
135            }
136        }
137    
138        protected void flushNextWriteBuffer() {
139            int nextnextSizesize = Math.min(Math.max(nextWriteBuffer.position(), 80), writeBufferSize);
140            ByteBuffer bb = nextWriteBuffer.toBuffer().toByteBuffer();
141            writeBuffer.add(bb);
142            writeBufferRemaining += bb.remaining();
143            nextWriteBuffer = new DataByteArrayOutputStream(nextnextSizesize);
144        }
145    
146        public ProtocolCodec.BufferState flush() throws IOException {
147            while (true) {
148                if (writeBufferRemaining != 0) {
149                    if( writeBuffer.size() == 1) {
150                        ByteBuffer b = writeBuffer.getFirst();
151                        lastWriteIoSize = writeChannel.write(b);
152                        if (lastWriteIoSize == 0) {
153                            return ProtocolCodec.BufferState.NOT_EMPTY;
154                        } else {
155                            writeBufferRemaining -= lastWriteIoSize;
156                            writeCounter += lastWriteIoSize;
157                            if(!b.hasRemaining()) {
158                                onBufferFlushed(writeBuffer.removeFirst());
159                            }
160                        }
161                    } else {
162                        ByteBuffer[] buffers = writeBuffer.toArray(new ByteBuffer[writeBuffer.size()]);
163                        lastWriteIoSize = writeChannel.write(buffers, 0, buffers.length);
164                        if (lastWriteIoSize == 0) {
165                            return ProtocolCodec.BufferState.NOT_EMPTY;
166                        } else {
167                            writeBufferRemaining -= lastWriteIoSize;
168                            writeCounter += lastWriteIoSize;
169                            while (!writeBuffer.isEmpty() && !writeBuffer.getFirst().hasRemaining()) {
170                                onBufferFlushed(writeBuffer.removeFirst());
171                            }
172                        }
173                    }
174                } else {
175                    if (nextWriteBuffer.size() == 0) {
176                        return ProtocolCodec.BufferState.EMPTY;
177                    } else {
178                        flushNextWriteBuffer();
179                    }
180                }
181            }
182        }
183    
184        /**
185         * Called when a buffer is flushed out.  Subclasses can implement
186         * in case they want to recycle the buffer.
187         *
188         * @param byteBuffer
189         */
190        protected void onBufferFlushed(ByteBuffer byteBuffer) {
191        }
192    
193        /////////////////////////////////////////////////////////////////////
194        //
195        // Non blocking read impl
196        //
197        /////////////////////////////////////////////////////////////////////
198    
199        abstract protected Action initialDecodeAction();
200    
201    
202        public void setReadableByteChannel(ReadableByteChannel channel) throws SocketException {
203            this.readChannel = channel;
204            if (this.readChannel instanceof SocketChannel) {
205                readBufferSize = ((SocketChannel) this.readChannel).socket().getReceiveBufferSize();
206            } else if (this.readChannel instanceof SslTransport.SSLChannel) {
207                writeBufferSize = ((SslTransport.SSLChannel) this.readChannel).socket().getReceiveBufferSize();
208            }
209            if( nextDecodeAction==null ) {
210                nextDecodeAction = initialDecodeAction();
211            }
212        }
213    
214        public void unread(byte[] buffer) {
215            assert ((readCounter == 0));
216            readBuffer.put(buffer);
217            readCounter += buffer.length;
218        }
219    
220        public long getReadCounter() {
221            return readCounter;
222        }
223    
224        public long getLastReadSize() {
225            return lastReadIoSize;
226        }
227    
228        public Object read() throws IOException {
229            Object command = null;
230            while (command == null) {
231                if (directReadBuffer != null) {
232                    while (directReadBuffer.hasRemaining()) {
233                        lastReadIoSize = readChannel.read(directReadBuffer);
234                        readCounter += lastReadIoSize;
235                        if (lastReadIoSize == -1) {
236                            throw new EOFException("Peer disconnected");
237                        } else if (lastReadIoSize == 0) {
238                            return null;
239                        }
240                    }
241                    command = nextDecodeAction.apply();
242                } else {
243                    if (readEnd == readBuffer.position()) {
244    
245                        if (readBuffer.remaining() == 0) {
246                            int size = readEnd - readStart;
247                            int newCapacity = 0;
248                            if (readStart == 0) {
249                                newCapacity = size + readBufferSize;
250                            } else {
251                                if (size > readBufferSize) {
252                                    newCapacity = size + readBufferSize;
253                                } else {
254                                    newCapacity = readBufferSize;
255                                }
256                            }
257                            byte[] newnewBufferbuffer = new byte[newCapacity];
258                            if (size > 0) {
259                                System.arraycopy(readBuffer.array(), readStart, newnewBufferbuffer, 0, size);
260                            }
261                            readBuffer = ByteBuffer.wrap(newnewBufferbuffer);
262                            readBuffer.position(size);
263                            readStart = 0;
264                            readEnd = size;
265                        }
266                        int p = readBuffer.position();
267                        lastReadIoSize = readChannel.read(readBuffer);
268                        readCounter += lastReadIoSize;
269                        if (lastReadIoSize == -1) {
270                            readCounter += 1; // to compensate for that -1
271                            throw new EOFException("Peer disconnected");
272                        } else if (lastReadIoSize == 0) {
273                            return null;
274                        }
275                    }
276                    command = nextDecodeAction.apply();
277                    assert ((readStart <= readEnd));
278                    assert ((readEnd <= readBuffer.position()));
279                }
280            }
281            return command;
282        }
283    
284        protected Buffer readUntil(Byte octet) throws ProtocolException {
285            return readUntil(octet, -1);
286        }
287    
288        protected Buffer readUntil(Byte octet, int max) throws ProtocolException {
289            return readUntil(octet, max, "Maximum protocol buffer length exeeded");
290        }
291    
292        protected Buffer readUntil(Byte octet, int max, String msg) throws ProtocolException {
293            byte[] array = readBuffer.array();
294            Buffer buf = new Buffer(array, readEnd, readBuffer.position() - readEnd);
295            int pos = buf.indexOf(octet);
296            if (pos >= 0) {
297                int offset = readStart;
298                readEnd += pos + 1;
299                readStart = readEnd;
300                int length = readEnd - offset;
301                if (max >= 0 && length > max) {
302                    throw new ProtocolException(msg);
303                }
304                return new Buffer(array, offset, length);
305            } else {
306                readEnd += buf.length;
307                if (max >= 0 && (readEnd - readStart) > max) {
308                    throw new ProtocolException(msg);
309                }
310                return null;
311            }
312        }
313    
314        protected Buffer readBytes(int length) {
315            if ((readBuffer.position() - readStart) < length) {
316                readEnd = readBuffer.position();
317                return null;
318            } else {
319                int offset = readStart;
320                readEnd = offset + length;
321                readStart = readEnd;
322                return new Buffer(readBuffer.array(), offset, length);
323            }
324        }
325    
326        protected Buffer peekBytes(int length) {
327            if ((readBuffer.position() - readStart) < length) {
328                readEnd = readBuffer.position();
329                return null;
330            } else {
331                return new Buffer(readBuffer.array(), readStart, length);
332            }
333        }
334    
335        protected Boolean readDirect(ByteBuffer buffer) {
336            assert (directReadBuffer == null || (directReadBuffer == buffer));
337    
338            if (buffer.hasRemaining()) {
339                // First we need to transfer the read bytes from the non-direct
340                // byte buffer into the direct one..
341                int limit = readBuffer.position();
342                int transferSize = Math.min((limit - readStart), buffer.remaining());
343                byte[] readBufferArray = readBuffer.array();
344                buffer.put(readBufferArray, readStart, transferSize);
345    
346                // The direct byte buffer might have been smaller than our readBuffer one..
347                // compact the readBuffer to avoid doing additional mem allocations.
348                int trailingSize = limit - (readStart + transferSize);
349                if (trailingSize > 0) {
350                    System.arraycopy(readBufferArray, readStart + transferSize, readBufferArray, readStart, trailingSize);
351                }
352                readBuffer.position(readStart + trailingSize);
353            }
354    
355            // For big direct byte buffers, it will still not have been filled,
356            // so install it so that we directly read into it until it is filled.
357            if (buffer.hasRemaining()) {
358                directReadBuffer = buffer;
359                return false;
360            } else {
361                directReadBuffer = null;
362                buffer.flip();
363                return true;
364            }
365        }
366    }