src/java/org/apache/cassandra/thrift/CustomTHsHaServer.java

package org.apache.cassandra.thrift;
/*
 * 
 * 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.
 * 
 */


import java.io.IOException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.spi.SelectorProvider;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Set;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.RejectedExecutionException;

import org.apache.cassandra.service.SocketSessionManagementService;
import org.apache.thrift.server.TNonblockingServer;
import org.apache.thrift.transport.TNonblockingServerTransport;
import org.apache.thrift.transport.TNonblockingSocket;
import org.apache.thrift.transport.TNonblockingTransport;
import org.apache.thrift.transport.TTransportException;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * This is a interim solution till THRIFT-1167 gets committed...
 * 
 * The idea here is to avoid sticking to one CPU for IO's. For better throughput
 * it is spread across multiple threads. Number of selector thread can be the
 * number of CPU available.
 */
public class CustomTHsHaServer extends TNonblockingServer
{
    private static final Logger LOGGER = LoggerFactory.getLogger(CustomTHsHaServer.class.getName());
    private Set<SelectorThread> ioThreads = new HashSet<SelectorThread>();
    private volatile boolean stopped_ = true;
    private ExecutorService invoker;

    /**
     * All the arguments to Non Blocking Server will apply here. In addition,
     * executor pool will be responsible for creating the internal threads which
     * will process the data. threads for selection usually are equal to the
     * number of cpu's
     */
    public CustomTHsHaServer(Args args, ExecutorService invoker, int threadCount)
    {
        super(args);
        this.invoker = invoker;
        // Create all the Network IO Threads.
        for (int i = 0; i < threadCount; ++i)
            ioThreads.add(new SelectorThread("Selector-Thread-" + i));
    }

    /** @inheritDoc */
    @Override
    public void serve()
    {
        if (!startListening())
            return;
        if (!startThreads())
            return;
        setServing(true);
        joinSelector();
        invoker.shutdown();
        setServing(false);
        stopListening();
    }

    /**
     * Save the remote socket as a thead local for future use of client state.
     */
    protected class Invocation implements Runnable
    {
        private final FrameBuffer frameBuffer;
        private SelectorThread thread;

        public Invocation(final FrameBuffer frameBuffer, SelectorThread thread)
        {
            this.frameBuffer = frameBuffer;
            this.thread = thread;
        }

        public void run()
        {
            TNonblockingSocket socket = (TNonblockingSocket) frameBuffer.trans_;
            SocketSessionManagementService.remoteSocket.set(socket.getSocketChannel().socket().getRemoteSocketAddress());
            frameBuffer.invoke();
            // this is how we let the same selector thread change the selection type.
            thread.requestSelectInterestChange(frameBuffer);
        }
    }

    protected boolean startThreads()
    {
        stopped_ = false;
        // start all the threads.
        for (SelectorThread thread : ioThreads)
            thread.start();
        return true;
    }

    @Override
    protected void joinSelector()
    {
        try
        {
            // wait till all done with stuff's
            for (SelectorThread thread : ioThreads)
                thread.join();
        } 
        catch (InterruptedException e)
        {
            LOGGER.error("Interrupted while joining threads!", e);
        }
    }

    /**
     * Stop serving and shut everything down.
     */
    @Override
    public void stop()
    {
        stopListening();
        stopped_ = true;
        for (SelectorThread thread : ioThreads)
            thread.wakeupSelector();
        joinSelector();
    }

    /**
     * IO Threads will perform expensive IO operations...
     */
    protected class SelectorThread extends Thread
    {
        private final Selector selector;
        private TNonblockingServerTransport serverTransport;
        private Set<FrameBuffer> selectInterestChanges = new HashSet<FrameBuffer>();

        public SelectorThread(String name)
        {
            super(name);
            try
            {
                this.selector = SelectorProvider.provider().openSelector();
                this.serverTransport = (TNonblockingServerTransport) serverTransport_;
                this.serverTransport.registerSelector(selector);
            } 
            catch (IOException ex)
            {
                throw new RuntimeException("Couldnt open the NIO selector", ex);
            }
        }

        public void run()
        {
            try
            {
                while (!stopped_)
                {
                    select();
                }
            } 
            catch (Throwable t)
            {
                LOGGER.error("Uncaught Exception: ", t);
            }
        }

        private void select() throws InterruptedException, IOException
        {
            // wait for new keys
            selector.select();
            Iterator<SelectionKey> keyIterator = selector.selectedKeys().iterator();
            while (keyIterator.hasNext())
            {
                SelectionKey key = keyIterator.next();
                keyIterator.remove();
                if (!key.isValid())
                {
                    // if invalid cleanup.
                    cleanupSelectionkey(key);
                    continue;
                }

                if (key.isAcceptable())
                    handleAccept();
                if (key.isReadable())
                    handleRead(key);
                else if (key.isWritable())
                    handleWrite(key);
                else
                    LOGGER.debug("Unexpected state " + key.interestOps());
            }
            // process the changes which are inserted after completion.
            processInterestChanges();
        }
        
        private void handleAccept()
        {
            SelectionKey clientKey = null;
            TNonblockingTransport client = null;
            try
            {
                // accept the connection
                client = (TNonblockingTransport) serverTransport.accept();
                clientKey = client.registerSelector(selector, SelectionKey.OP_READ);
                // add this key to the map
                FrameBuffer frameBuffer = new FrameBuffer(client, clientKey);
                clientKey.attach(frameBuffer); 
            } catch (TTransportException ex)
            {
                // ignore this might have been handled by the other threads.
                // serverTransport.accept() as it returns null as nothing to accept.
                return;
            }
            catch (IOException tte)
            {
                // something went wrong accepting.
                LOGGER.warn("Exception trying to accept!", tte);
                tte.printStackTrace();
                if (clientKey != null)
                    cleanupSelectionkey(clientKey);
                if (client != null)
                    client.close();
            }
        }
        
        private void handleRead(SelectionKey key)
        {
            FrameBuffer buffer = (FrameBuffer) key.attachment();
            if (!buffer.read())
            {
                cleanupSelectionkey(key);
                return;
            }

            if (buffer.isFrameFullyRead())
            {
                if (!requestInvoke(buffer, this))
                    cleanupSelectionkey(key);
            }
        }
        
        private void handleWrite(SelectionKey key)
        {
            FrameBuffer buffer = (FrameBuffer) key.attachment();
            if (!buffer.write())
                cleanupSelectionkey(key);
        }
        
        public void requestSelectInterestChange(FrameBuffer frameBuffer)
        {
            synchronized (selectInterestChanges)
            {
                selectInterestChanges.add(frameBuffer);
            }
            // Wake-up the selector, if it's currently blocked.
            selector.wakeup();
        }

        private void processInterestChanges()
        {
            synchronized (selectInterestChanges)
            {
                for (FrameBuffer fb : selectInterestChanges)
                    fb.changeSelectInterests();
                selectInterestChanges.clear();
            }
        }
        
        private void cleanupSelectionkey(SelectionKey key)
        {
            FrameBuffer buffer = (FrameBuffer) key.attachment();
            if (buffer != null)
                buffer.close();
            // cancel the selection key
            key.cancel();
        }
        
        public void wakeupSelector()
        {
            selector.wakeup();
        }
    }
    
    protected boolean requestInvoke(FrameBuffer frameBuffer, SelectorThread thread)
    {
        try
        {
            Runnable invocation = new Invocation(frameBuffer, thread);
            invoker.execute(invocation);
            return true;
        } 
        catch (RejectedExecutionException rx)
        {
            LOGGER.warn("ExecutorService rejected execution!", rx);
            return false;
        }
    }

    @Override
    protected void requestSelectInterestChange(FrameBuffer fb)
    {
        // Dont change the interest here, this has to be done by the selector
        // thread because the method is not synchronized with the rest of the
        // selectors threads.
    }
}