001    /**
002     * Copyright (C) 2010-2011, FuseSource Corp.  All rights reserved.
003     *
004     *     http://fusesource.com
005     *
006     * The software in this package is published under the terms of the
007     * CDDL license a copy of which has been included with this distribution
008     * in the license.txt file.
009     */
010    package org.fusesource.hawtdispatch.transport;
011    
012    import org.fusesource.hawtdispatch.*;
013    
014    import java.io.IOException;
015    import java.net.*;
016    import java.nio.ByteBuffer;
017    import java.nio.channels.ReadableByteChannel;
018    import java.nio.channels.SelectionKey;
019    import java.nio.channels.SocketChannel;
020    import java.nio.channels.WritableByteChannel;
021    import java.util.LinkedList;
022    import java.util.concurrent.TimeUnit;
023    
024    /**
025     * An implementation of the {@link org.fusesource.hawtdispatch.transport.Transport} interface using raw tcp/ip
026     *
027     * @author <a href="http://hiramchirino.com">Hiram Chirino</a>
028     */
029    public class TcpTransport extends ServiceBase implements Transport {
030    
031        abstract static class SocketState {
032            void onStop(Runnable onCompleted) {
033            }
034            void onCanceled() {
035            }
036            boolean is(Class<? extends SocketState> clazz) {
037                return getClass()==clazz;
038            }
039        }
040    
041        static class DISCONNECTED extends SocketState{}
042    
043        class CONNECTING extends SocketState{
044            void onStop(Runnable onCompleted) {
045                trace("CONNECTING.onStop");
046                CANCELING state = new CANCELING();
047                socketState = state;
048                state.onStop(onCompleted);
049            }
050            void onCanceled() {
051                trace("CONNECTING.onCanceled");
052                CANCELING state = new CANCELING();
053                socketState = state;
054                state.onCanceled();
055            }
056        }
057    
058        class CONNECTED extends SocketState {
059    
060            public CONNECTED() {
061                localAddress = channel.socket().getLocalSocketAddress();
062                remoteAddress = channel.socket().getRemoteSocketAddress();
063            }
064    
065            void onStop(Runnable onCompleted) {
066                trace("CONNECTED.onStop");
067                CANCELING state = new CANCELING();
068                socketState = state;
069                state.add(createDisconnectTask());
070                state.onStop(onCompleted);
071            }
072            void onCanceled() {
073                trace("CONNECTED.onCanceled");
074                CANCELING state = new CANCELING();
075                socketState = state;
076                state.add(createDisconnectTask());
077                state.onCanceled();
078            }
079            Runnable createDisconnectTask() {
080                return new Runnable(){
081                    public void run() {
082                        listener.onTransportDisconnected(false);
083                    }
084                };
085            }
086        }
087    
088        class CANCELING extends SocketState {
089            private LinkedList<Runnable> runnables =  new LinkedList<Runnable>();
090            private int remaining;
091            private boolean dispose;
092    
093            public CANCELING() {
094                if( readSource!=null ) {
095                    remaining++;
096                    readSource.cancel();
097                }
098                if( writeSource!=null ) {
099                    remaining++;
100                    writeSource.cancel();
101                }
102            }
103            void onStop(Runnable onCompleted) {
104                trace("CANCELING.onCompleted");
105                add(onCompleted);
106                dispose = true;
107            }
108            void add(Runnable onCompleted) {
109                if( onCompleted!=null ) {
110                    runnables.add(onCompleted);
111                }
112            }
113            void onCanceled() {
114                trace("CANCELING.onCanceled");
115                remaining--;
116                if( remaining!=0 ) {
117                    return;
118                }
119                try {
120                    channel.close();
121                } catch (IOException ignore) {
122                }
123                socketState = new CANCELED(dispose);
124                for (Runnable runnable : runnables) {
125                    runnable.run();
126                }
127                if (dispose) {
128                    dispose();
129                }
130            }
131        }
132    
133        class CANCELED extends SocketState {
134            private boolean disposed;
135    
136            public CANCELED(boolean disposed) {
137                this.disposed=disposed;
138            }
139    
140            void onStop(Runnable onCompleted) {
141                trace("CANCELED.onStop");
142                if( !disposed ) {
143                    disposed = true;
144                    dispose();
145                }
146                onCompleted.run();
147            }
148        }
149    
150        protected URI remoteLocation;
151        protected URI localLocation;
152        protected TransportListener listener;
153        protected ProtocolCodec codec;
154    
155        protected SocketChannel channel;
156    
157        protected SocketState socketState = new DISCONNECTED();
158    
159        protected DispatchQueue dispatchQueue;
160        private DispatchSource readSource;
161        private DispatchSource writeSource;
162        protected CustomDispatchSource<Integer, Integer> drainOutboundSource;
163        protected CustomDispatchSource<Integer, Integer> yieldSource;
164    
165        protected boolean useLocalHost = true;
166    
167        int maxReadRate;
168        int maxWriteRate;
169        int receiveBufferSize = 1024*64;
170        int sendBufferSize = 1024*64;
171        boolean keepAlive = true;
172    
173    
174        public static final int IPTOS_LOWCOST = 0x02;
175        public static final int IPTOS_RELIABILITY = 0x04;
176        public static final int IPTOS_THROUGHPUT = 0x08;
177        public static final int IPTOS_LOWDELAY = 0x10;
178    
179        int trafficClass = IPTOS_THROUGHPUT;
180    
181        protected RateLimitingChannel rateLimitingChannel;
182        SocketAddress localAddress;
183        SocketAddress remoteAddress;
184    
185        class RateLimitingChannel implements ReadableByteChannel, WritableByteChannel {
186    
187            int read_allowance = maxReadRate;
188            boolean read_suspended = false;
189            int read_resume_counter = 0;
190            int write_allowance = maxWriteRate;
191            boolean write_suspended = false;
192    
193            public void resetAllowance() {
194                if( read_allowance != maxReadRate || write_allowance != maxWriteRate) {
195                    read_allowance = maxReadRate;
196                    write_allowance = maxWriteRate;
197                    if( write_suspended ) {
198                        write_suspended = false;
199                        resumeWrite();
200                    }
201                    if( read_suspended ) {
202                        read_suspended = false;
203                        resumeRead();
204                        for( int i=0; i < read_resume_counter ; i++ ) {
205                            resumeRead();
206                        }
207                    }
208                }
209            }
210    
211            public int read(ByteBuffer dst) throws IOException {
212                if( maxReadRate ==0 ) {
213                    return channel.read(dst);
214                } else {
215                    int remaining = dst.remaining();
216                    if( read_allowance ==0 || remaining ==0 ) {
217                        return 0;
218                    }
219    
220                    int reduction = 0;
221                    if( remaining > read_allowance) {
222                        reduction = remaining - read_allowance;
223                        dst.limit(dst.limit() - reduction);
224                    }
225                    int rc=0;
226                    try {
227                        rc = channel.read(dst);
228                        read_allowance -= rc;
229                    } finally {
230                        if( reduction!=0 ) {
231                            if( dst.remaining() == 0 ) {
232                                // we need to suspend the read now until we get
233                                // a new allowance..
234                                readSource.suspend();
235                                read_suspended = true;
236                            }
237                            dst.limit(dst.limit() + reduction);
238                        }
239                    }
240                    return rc;
241                }
242            }
243    
244            public int write(ByteBuffer src) throws IOException {
245                if( maxWriteRate ==0 ) {
246                    return channel.write(src);
247                } else {
248                    int remaining = src.remaining();
249                    if( write_allowance ==0 || remaining ==0 ) {
250                        return 0;
251                    }
252    
253                    int reduction = 0;
254                    if( remaining > write_allowance) {
255                        reduction = remaining - write_allowance;
256                        src.limit(src.limit() - reduction);
257                    }
258                    int rc = 0;
259                    try {
260                        rc = channel.write(src);
261                        write_allowance -= rc;
262                    } finally {
263                        if( reduction!=0 ) {
264                            if( src.remaining() == 0 ) {
265                                // we need to suspend the read now until we get
266                                // a new allowance..
267                                write_suspended = true;
268                                suspendWrite();
269                            }
270                            src.limit(src.limit() + reduction);
271                        }
272                    }
273                    return rc;
274                }
275            }
276    
277            public boolean isOpen() {
278                return channel.isOpen();
279            }
280    
281            public void close() throws IOException {
282                channel.close();
283            }
284    
285            public void resumeRead() {
286                if( read_suspended ) {
287                    read_resume_counter += 1;
288                } else {
289                    _resumeRead();
290                }
291            }
292    
293        }
294    
295        private final Runnable CANCEL_HANDLER = new Runnable() {
296            public void run() {
297                socketState.onCanceled();
298            }
299        };
300    
301        static final class OneWay {
302            final Object command;
303            final Retained retained;
304    
305            public OneWay(Object command, Retained retained) {
306                this.command = command;
307                this.retained = retained;
308            }
309        }
310    
311        public void connected(SocketChannel channel) throws IOException, Exception {
312            this.channel = channel;
313            initializeChannel();
314            this.socketState = new CONNECTED();
315        }
316    
317        protected void initializeChannel() throws Exception {
318            this.channel.configureBlocking(false);
319            Socket socket = channel.socket();
320            try {
321                socket.setReuseAddress(true);
322            } catch (SocketException e) {
323            }
324            try {
325                socket.setSoLinger(true, 0);
326            } catch (SocketException e) {
327            }
328            try {
329                socket.setTrafficClass(trafficClass);
330            } catch (SocketException e) {
331            }
332            try {
333                socket.setKeepAlive(keepAlive);
334            } catch (SocketException e) {
335            }
336            try {
337                socket.setTcpNoDelay(true);
338            } catch (SocketException e) {
339            }
340            try {
341                socket.setReceiveBufferSize(receiveBufferSize);
342            } catch (SocketException e) {
343            }
344            try {
345                socket.setSendBufferSize(sendBufferSize);
346            } catch (SocketException e) {
347            }
348    
349            if( channel!=null && codec!=null ) {
350                initializeCodec();
351            }
352        }
353    
354        protected void initializeCodec() throws Exception {
355            codec.setReadableByteChannel(readChannel());
356            codec.setWritableByteChannel(writeChannel());
357        }
358    
359        public void connecting(URI remoteLocation, URI localLocation) throws IOException, Exception {
360            this.channel = SocketChannel.open();
361            initializeChannel();
362            this.remoteLocation = remoteLocation;
363            this.localLocation = localLocation;
364    
365            if (localLocation != null) {
366                InetSocketAddress localAddress = new InetSocketAddress(InetAddress.getByName(localLocation.getHost()), localLocation.getPort());
367                channel.socket().bind(localAddress);
368            }
369    
370            String host = resolveHostName(remoteLocation.getHost());
371            InetSocketAddress remoteAddress = new InetSocketAddress(host, remoteLocation.getPort());
372            channel.connect(remoteAddress);
373            this.socketState = new CONNECTING();
374        }
375    
376    
377        public DispatchQueue getDispatchQueue() {
378            return dispatchQueue;
379        }
380    
381        public void setDispatchQueue(DispatchQueue queue) {
382            this.dispatchQueue = queue;
383        }
384    
385        public void _start(Runnable onCompleted) {
386            try {
387                if (socketState.is(CONNECTING.class) ) {
388                    trace("connecting...");
389                    // this allows the connect to complete..
390                    readSource = Dispatch.createSource(channel, SelectionKey.OP_CONNECT, dispatchQueue);
391                    readSource.setEventHandler(new Runnable() {
392                        public void run() {
393                            if (getServiceState() != STARTED) {
394                                return;
395                            }
396                            try {
397                                trace("connected.");
398                                channel.finishConnect();
399                                readSource.setCancelHandler(null);
400                                readSource.cancel();
401                                readSource=null;
402                                socketState = new CONNECTED();
403                                onConnected();
404                            } catch (IOException e) {
405                                onTransportFailure(e);
406                            }
407                        }
408                    });
409                    readSource.setCancelHandler(CANCEL_HANDLER);
410                    readSource.resume();
411    
412                } else if (socketState.is(CONNECTED.class) ) {
413                    dispatchQueue.execute(new Runnable() {
414                        public void run() {
415                            try {
416                                trace("was connected.");
417                                onConnected();
418                            } catch (IOException e) {
419                                 onTransportFailure(e);
420                            }
421                        }
422                    });
423                } else {
424                    System.err.println("cannot be started.  socket state is: "+socketState);
425                }
426            } finally {
427                if( onCompleted!=null ) {
428                    onCompleted.run();
429                }
430            }
431        }
432    
433        public void _stop(final Runnable onCompleted) {
434            trace("stopping.. at state: "+socketState);
435            socketState.onStop(onCompleted);
436        }
437    
438        protected String resolveHostName(String host) throws UnknownHostException {
439            String localName = InetAddress.getLocalHost().getHostName();
440            if (localName != null && isUseLocalHost()) {
441                if (localName.equals(host)) {
442                    return "localhost";
443                }
444            }
445            return host;
446        }
447    
448        protected void onConnected() throws IOException {
449            yieldSource = Dispatch.createSource(EventAggregators.INTEGER_ADD, dispatchQueue);
450            yieldSource.setEventHandler(new Runnable() {
451                public void run() {
452                    drainInbound();
453                }
454            });
455            yieldSource.resume();
456            drainOutboundSource = Dispatch.createSource(EventAggregators.INTEGER_ADD, dispatchQueue);
457            drainOutboundSource.setEventHandler(new Runnable() {
458                public void run() {
459                    flush();
460                }
461            });
462            drainOutboundSource.resume();
463    
464            readSource = Dispatch.createSource(channel, SelectionKey.OP_READ, dispatchQueue);
465            writeSource = Dispatch.createSource(channel, SelectionKey.OP_WRITE, dispatchQueue);
466    
467            readSource.setCancelHandler(CANCEL_HANDLER);
468            writeSource.setCancelHandler(CANCEL_HANDLER);
469    
470            readSource.setEventHandler(new Runnable() {
471                public void run() {
472                    drainInbound();
473                }
474            });
475            writeSource.setEventHandler(new Runnable() {
476                public void run() {
477                    flush();
478                }
479            });
480    
481            if( maxReadRate !=0 || maxWriteRate !=0 ) {
482                rateLimitingChannel = new RateLimitingChannel();
483                schedualRateAllowanceReset();
484            }
485            listener.onTransportConnected();
486        }
487    
488        private void schedualRateAllowanceReset() {
489            dispatchQueue.executeAfter(1, TimeUnit.SECONDS, new Runnable(){
490                public void run() {
491                    if( !socketState.is(CONNECTED.class) ) {
492                        return;
493                    }
494                    rateLimitingChannel.resetAllowance();
495                    schedualRateAllowanceReset();
496                }
497            });
498        }
499    
500        private void dispose() {
501            if( readSource!=null ) {
502                readSource.cancel();
503                readSource=null;
504            }
505    
506            if( writeSource!=null ) {
507                writeSource.cancel();
508                writeSource=null;
509            }
510            this.codec = null;
511        }
512    
513        public void onTransportFailure(IOException error) {
514            listener.onTransportFailure(error);
515            socketState.onCanceled();
516        }
517    
518    
519        public boolean full() {
520            return codec==null || codec.full();
521        }
522    
523        boolean rejectingOffers;
524    
525        public boolean offer(Object command) {
526            dispatchQueue.assertExecuting();
527            try {
528                if (!socketState.is(CONNECTED.class)) {
529                    throw new IOException("Not connected.");
530                }
531                if (getServiceState() != STARTED) {
532                    throw new IOException("Not running.");
533                }
534    
535                ProtocolCodec.BufferState rc = codec.write(command);
536                rejectingOffers = codec.full();
537                switch (rc ) {
538                    case FULL:
539                        return false;
540                    default:
541                        drainOutboundSource.merge(1);
542                        return true;
543                }
544            } catch (IOException e) {
545                onTransportFailure(e);
546                return false;
547            }
548    
549        }
550    
551        boolean writeResumedForCodecFlush = false;
552    
553        /**
554         *
555         */
556        public void flush() {
557            dispatchQueue.assertExecuting();
558            if (getServiceState() != STARTED || !socketState.is(CONNECTED.class)) {
559                return;
560            }
561            try {
562                if( codec.flush() == ProtocolCodec.BufferState.EMPTY && transportFlush() ) {
563                    if( writeResumedForCodecFlush) {
564                        writeResumedForCodecFlush = false;
565                        suspendWrite();
566                    }
567                    rejectingOffers = false;
568                    listener.onRefill();
569    
570                } else {
571                    if(!writeResumedForCodecFlush) {
572                        writeResumedForCodecFlush = true;
573                        resumeWrite();
574                    }
575                }
576            } catch (IOException e) {
577                onTransportFailure(e);
578            }
579        }
580    
581        protected boolean transportFlush() throws IOException {
582            return true;
583        }
584    
585        protected void drainInbound() {
586            if (!getServiceState().isStarted() || readSource.isSuspended()) {
587                return;
588            }
589            try {
590                long initial = codec.getReadCounter();
591                // Only process upto 2 x the read buffer worth of data at a time so we can give
592                // other connections a chance to process their requests.
593                while( codec.getReadCounter()-initial < codec.getReadBufferSize()<<2 ) {
594                    Object command = codec.read();
595                    if ( command!=null ) {
596                        try {
597                            listener.onTransportCommand(command);
598                        } catch (Throwable e) {
599                            e.printStackTrace();
600                            onTransportFailure(new IOException("Transport listener failure."));
601                        }
602    
603                        // the transport may be suspended after processing a command.
604                        if (getServiceState() == STOPPED || readSource.isSuspended()) {
605                            return;
606                        }
607                    } else {
608                        return;
609                    }
610                }
611                yieldSource.merge(1);
612            } catch (IOException e) {
613                onTransportFailure(e);
614            }
615        }
616    
617        public SocketAddress getLocalAddress() {
618            return localAddress;
619        }
620    
621        public SocketAddress getRemoteAddress() {
622            return remoteAddress;
623        }
624    
625        public <T> T narrow(Class<T> target) {
626            if (target.isAssignableFrom(getClass())) {
627                return target.cast(this);
628            }
629            return null;
630        }
631    
632        private boolean assertConnected() {
633            try {
634                if ( !isConnected() ) {
635                    throw new IOException("Not connected.");
636                }
637                return true;
638            } catch (IOException e) {
639                onTransportFailure(e);
640            }
641            return false;
642        }
643    
644        public void suspendRead() {
645            if( isConnected() && readSource!=null ) {
646                readSource.suspend();
647            }
648        }
649    
650    
651        public void resumeRead() {
652            if( isConnected() && readSource!=null ) {
653                if( rateLimitingChannel!=null ) {
654                    rateLimitingChannel.resumeRead();
655                } else {
656                    _resumeRead();
657                }
658            }
659        }
660    
661        private void _resumeRead() {
662            readSource.resume();
663            dispatchQueue.execute(new Runnable(){
664                public void run() {
665                    drainInbound();
666                }
667            });
668        }
669    
670        protected void suspendWrite() {
671            if( isConnected() && writeSource!=null ) {
672                writeSource.suspend();
673            }
674        }
675    
676        protected void resumeWrite() {
677            if( isConnected() && writeSource!=null ) {
678                writeSource.resume();
679            }
680        }
681    
682        public String getTypeId() {
683            return "tcp";
684        }
685    
686        public void reconnect(URI uri) {
687            throw new UnsupportedOperationException();
688        }
689    
690        public TransportListener getTransportListener() {
691            return listener;
692        }
693    
694        public void setTransportListener(TransportListener listener) {
695            this.listener = listener;
696        }
697    
698        public ProtocolCodec getProtocolCodec() {
699            return codec;
700        }
701    
702        public void setProtocolCodec(ProtocolCodec protocolCodec) throws Exception {
703            this.codec = protocolCodec;
704            if( channel!=null && codec!=null ) {
705                initializeCodec();
706            }
707        }
708    
709        public boolean isConnected() {
710            return socketState.is(CONNECTED.class);
711        }
712    
713        public boolean isDisposed() {
714            return getServiceState() == STOPPED;
715        }
716    
717        public boolean isFaultTolerant() {
718            return false;
719        }
720    
721        public boolean isUseLocalHost() {
722            return useLocalHost;
723        }
724    
725        /**
726         * Sets whether 'localhost' or the actual local host name should be used to
727         * make local connections. On some operating systems such as Macs its not
728         * possible to connect as the local host name so localhost is better.
729         */
730        public void setUseLocalHost(boolean useLocalHost) {
731            this.useLocalHost = useLocalHost;
732        }
733    
734        private void trace(String message) {
735            // TODO:
736        }
737    
738        public SocketChannel getSocketChannel() {
739            return channel;
740        }
741    
742        public ReadableByteChannel readChannel() {
743            if(rateLimitingChannel!=null) {
744                return rateLimitingChannel;
745            } else {
746                return channel;
747            }
748        }
749    
750        public WritableByteChannel writeChannel() {
751            if(rateLimitingChannel!=null) {
752                return rateLimitingChannel;
753            } else {
754                return channel;
755            }
756        }
757    
758        public int getMaxReadRate() {
759            return maxReadRate;
760        }
761    
762        public void setMaxReadRate(int maxReadRate) {
763            this.maxReadRate = maxReadRate;
764        }
765    
766        public int getMaxWriteRate() {
767            return maxWriteRate;
768        }
769    
770        public void setMaxWriteRate(int maxWriteRate) {
771            this.maxWriteRate = maxWriteRate;
772        }
773    
774        public int getTrafficClass() {
775            return trafficClass;
776        }
777    
778        public void setTrafficClass(int trafficClass) {
779            this.trafficClass = trafficClass;
780        }
781    
782        public int getReceiveBufferSize() {
783            return receiveBufferSize;
784        }
785    
786        public void setReceiveBufferSize(int receiveBufferSize) {
787            this.receiveBufferSize = receiveBufferSize;
788        }
789    
790        public int getSendBufferSize() {
791            return sendBufferSize;
792        }
793    
794        public void setSendBufferSize(int sendBufferSize) {
795            this.sendBufferSize = sendBufferSize;
796        }
797    
798        public boolean isKeepAlive() {
799            return keepAlive;
800        }
801    
802        public void setKeepAlive(boolean keepAlive) {
803            this.keepAlive = keepAlive;
804        }
805    
806    }