Monday, November 5, 2012

How to stop biting, when you cant chew more..

This is a follow up to my earlier post 'Does Tomcat bite more than it can chew?' and illustrates a pure Java program that utilizes Java NIO to stop accepting new messages when one is not able to handle the load, without any dependence on the TCP backlog etc.

Program Implementation

We open a selector, and invoke the startListening() method, that opens a ServerSocketChannel and then binds it to port 8280. The channel is configured as non-blocking, and finally we register our interest in OP_ACCEPT to handle incoming connections.

    private void startListening() throws IOException {
        server = ServerSocketChannel.open();
        server.socket().bind(new InetSocketAddress(8280), 0);
        server.configureBlocking(false);
        server.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("\nI am ready to listen for new messages now..");
    }


If you telnet to the port via a command line after the server starts up, you would see a message "Hi there! type a word". The server accepts the incoming connection as a non-blocking connection, and registers OP_READ to read the content typed in.


To illustrate how the server can prevent another client from connecting to it while it serves the currently connected client, it prints "I accepted this one.. but not any more now" on its console, cancels the SelectionKey and closes the channel.

A new telnet session will see the "Connection refused" error as expected.

Next, I would type a small word into the first telnet session, and the server would print it in its console, and close the connection.

At the same time, it prints the message "I am ready to listen for new messages now.." on its console, and invokes the above startListening() method again - making it ready to listen and accept a new client.

We re-try the connection from the same command prompt that received the 'Connection refused' earlier, and as expected are greeted with the welcome message again.

What did we learn?

A low level NIO server can stop accepting new connections, if it can determine that its not able to serve new clients. Once it stops accepting connections this way, any client connection attempt will see a 'Connection Refused' error. This may not be the case if our server was implemented differently (See my last article and its example and how Tomcat behaves under load).

Complete source code

import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.nio.channels.ServerSocketChannel;
import java.nio.channels.SocketChannel;
import java.util.Iterator;

public class TestAccept2 {

    private ServerSocketChannel server = null;
    private Selector selector = null;

    public static void main(String[] args) throws Exception {
        new TestAccept2().run();
    }

    private void run() throws Exception {
        selector = Selector.open();
        startListening();

        while (true) {
            selector.select();

            for (Iterator i = selector.selectedKeys().iterator(); i.hasNext(); ) {
                SelectionKey key = i.next();
                i.remove();
                if (key.isAcceptable()) {
                    SocketChannel client = server.accept();
                    client.configureBlocking(false);
                    client.socket().setTcpNoDelay(true);
                    client.register(selector, SelectionKey.OP_READ);

                    System.out.println("I accepted this one.. but not any more now");
                    key.cancel();
                    key.channel().close();
                    sayHello(client);

                } else if (key.isReadable()) {
                    readDataFromSocket(key);
                }
            }
        }
    }

    private void startListening() throws IOException {
        server = ServerSocketChannel.open();
        server.socket().bind(new InetSocketAddress(8280), 0);
        server.configureBlocking(false);
        server.register(selector, SelectionKey.OP_ACCEPT);
        System.out.println("\nI am ready to listen for new messages now..");
    }

    private void sayHello(SocketChannel channel) throws Exception {
        channel.write(ByteBuffer.wrap("Hi there! type a word\r\n".getBytes()));
    }

    private void readDataFromSocket(SelectionKey key) throws Exception {
        SocketChannel socketChannel = (SocketChannel) key.channel();
        ByteBuffer buffer = ByteBuffer.allocate(32);
        if (socketChannel.read(buffer) > 0) {
            buffer.flip();
            byte[] bytearr = new byte[buffer.remaining()];
            buffer.get(bytearr);
            System.out.print(new String(bytearr));
            socketChannel.close();

            startListening();
        }
    }

}