src/com/gitblit/fanout/FanoutNioService.java - gitblit - Git at Google

 /*
  * Copyright 2013 gitblit.com.
  *
  * Licensed 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.
  */
 package com.gitblit.fanout;

 import java.io.IOException;
 import java.net.InetSocketAddress;
 import java.nio.ByteBuffer;
 import java.nio.CharBuffer;
 import java.nio.channels.SelectionKey;
 import java.nio.channels.Selector;
 import java.nio.channels.ServerSocketChannel;
 import java.nio.channels.SocketChannel;
 import java.nio.charset.CharacterCodingException;
 import java.nio.charset.Charset;
 import java.nio.charset.CharsetDecoder;
 import java.text.MessageFormat;
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.HashMap;
 import java.util.Iterator;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;

 import org.slf4j.Logger;
 import org.slf4j.LoggerFactory;

 /**
  * A single-thread NIO implementation of https://github.com/travisghansen/fanout
  *
  * This implementation uses channels and selectors, which are the Java analog of
  * the Linux epoll mechanism used in the original fanout C code.
  *
  * @author James Moger
  *
  */
 public class FanoutNioService extends FanoutService {

 	private final static Logger logger = LoggerFactory.getLogger(FanoutNioService.class);

 	private volatile ServerSocketChannel serviceCh;
 	private volatile Selector selector;

 	public static void main(String[] args) throws Exception {
 		FanoutNioService pubsub = new FanoutNioService(null, DEFAULT_PORT);
 		pubsub.setStrictRequestTermination(false);
 		pubsub.setAllowAllChannelAnnouncements(false);
 		pubsub.start();
 	}

 	/**
 	 * Create a single-threaded fanout service.
 	 *
 	 * @param host
 	 * @param port
 	 *            the port for running the fanout PubSub service
 	 * @throws IOException
 	 */
 	public FanoutNioService(int port) {
 		this(null, port);
 	}

 	/**
 	 * Create a single-threaded fanout service.
 	 *
 	 * @param bindInterface
 	 *            the ip address to bind for the service, may be null
 	 * @param port
 	 *            the port for running the fanout PubSub service
 	 * @throws IOException
 	 */
 	public FanoutNioService(String bindInterface, int port) {
 		super(bindInterface, port, "Fanout nio service");
 	}

 	@Override
 	protected boolean isConnected() {
 		return serviceCh != null;
 	}

 	@Override
 	protected boolean connect() {
 		if (serviceCh == null) {
 			try {
 				serviceCh = ServerSocketChannel.open();
 				serviceCh.configureBlocking(false);
 				serviceCh.socket().setReuseAddress(true);
 				serviceCh.socket().bind(host == null ? new InetSocketAddress(port) : new InetSocketAddress(host, port));
 				selector = Selector.open();
 				serviceCh.register(selector, SelectionKey.OP_ACCEPT);
 				logger.info(MessageFormat.format("{0} is ready on {1}:{2,number,0}",
 						name, host == null ? "0.0.0.0" : host, port));
 			} catch (IOException e) {
 				logger.error(MessageFormat.format("failed to open {0} on {1}:{2,number,0}",
 						name, name, host == null ? "0.0.0.0" : host, port), e);
 				return false;
 			}
 		}
 		return true;
 	}

 	@Override
 	protected void disconnect() {
 		try {
 			if (serviceCh != null) {
 				// close all active client connections
 				Map<String, SocketChannel> clients = getCurrentClientSockets();
 				for (Map.Entry<String, SocketChannel> client : clients.entrySet()) {
 					closeClientSocket(client.getKey(), client.getValue());
 				}

 				// close service socket channel
 				logger.debug(MessageFormat.format("closing {0} socket channel", name));
 				serviceCh.socket().close();
 				serviceCh.close();
 				serviceCh = null;
 				selector.close();
 				selector = null;
 			}
 		} catch (IOException e) {
 			logger.error(MessageFormat.format("failed to disconnect {0}", name), e);
 		}
 	}

 	@Override
 	protected void listen() throws IOException {
 		while (selector.select(serviceTimeout) > 0) {
 			Set<SelectionKey> keys = selector.selectedKeys();
 			Iterator<SelectionKey> keyItr = keys.iterator();
 			while (keyItr.hasNext()) {
 				SelectionKey key = (SelectionKey) keyItr.next();
 				if (key.isAcceptable()) {
 					// new fanout client connection
 					ServerSocketChannel sch = (ServerSocketChannel) key.channel();
 					try {
 						SocketChannel ch = sch.accept();
 						ch.configureBlocking(false);
 						configureClientSocket(ch.socket());

 						FanoutNioConnection connection = new FanoutNioConnection(ch);
 						addConnection(connection);

 						// register to send the queued message
 						ch.register(selector, SelectionKey.OP_WRITE, connection);
 					} catch (IOException e) {
 						logger.error("error accepting fanout connection", e);
 					}
 				} else if (key.isReadable()) {
 					// read fanout client request
 					SocketChannel ch = (SocketChannel) key.channel();
 					FanoutNioConnection connection = (FanoutNioConnection) key.attachment();
 					try {
 						connection.read(ch, isStrictRequestTermination());
 						int replies = 0;
 						Iterator<String> reqItr = connection.requestQueue.iterator();
 						while (reqItr.hasNext()) {
 							String req = reqItr.next();
 							String reply = processRequest(connection, req);
 							reqItr.remove();
 							if (reply != null) {
 								replies++;
 							}
 						}

 						if (replies > 0) {
 							// register to send the replies to requests
 							ch.register(selector, SelectionKey.OP_WRITE, connection);
 						} else {
 							// re-register for next read
 							ch.register(selector, SelectionKey.OP_READ, connection);
 						}
 					} catch (IOException e) {
 						logger.error(MessageFormat.format("fanout connection {0} error: {1}", connection.id, e.getMessage()));
 						removeConnection(connection);
 						closeClientSocket(connection.id, ch);
 					}
 				} else if (key.isWritable()) {
 					// asynchronous reply to fanout client request
 					SocketChannel ch = (SocketChannel) key.channel();
 					FanoutNioConnection connection = (FanoutNioConnection) key.attachment();
 					try {
 						connection.write(ch);

 						if (hasConnection(connection)) {
 							// register for next read
 							ch.register(selector, SelectionKey.OP_READ, connection);
 						} else {
 							// Connection was rejected due to load or
 							// some other reason. Close it.
 							closeClientSocket(connection.id, ch);
 						}
 					} catch (IOException e) {
 						logger.error(MessageFormat.format("fanout connection {0}: {1}", connection.id, e.getMessage()));
 						removeConnection(connection);
 						closeClientSocket(connection.id, ch);
 					}
 				}
 				keyItr.remove();
 			}
 		}
 	}

 	protected void closeClientSocket(String id, SocketChannel ch) {
 		try {
 			ch.close();
 		} catch (IOException e) {
 			logger.error(MessageFormat.format("fanout connection {0}", id), e);
 		}
 	}

 	protected void broadcast(Collection<FanoutServiceConnection> connections, String channel, String message) {
 		super.broadcast(connections, channel, message);

 		// register queued write
 		Map<String, SocketChannel> sockets = getCurrentClientSockets();
 		for (FanoutServiceConnection connection : connections) {
 			SocketChannel ch = sockets.get(connection.id);
 			if (ch == null) {
 				logger.warn(MessageFormat.format("fanout connection {0} has been disconnected", connection.id));
 				removeConnection(connection);
 				continue;
 			}
 			try {
 				ch.register(selector, SelectionKey.OP_WRITE, connection);
 			} catch (IOException e) {
 				logger.error(MessageFormat.format("failed to register write op for fanout connection {0}", connection.id));
 			}
 		}
 	}

 	protected Map<String, SocketChannel> getCurrentClientSockets() {
 		Map<String, SocketChannel> sockets = new HashMap<String, SocketChannel>();
 		for (SelectionKey key : selector.keys()) {
 			if (key.channel() instanceof SocketChannel) {
 				SocketChannel ch = (SocketChannel) key.channel();
 				String id = FanoutConstants.getRemoteSocketId(ch.socket());
 				sockets.put(id, ch);
 			}
 		}
 		return sockets;
 	}

 	/**
 	 * FanoutNioConnection handles reading/writing messages from a remote fanout
 	 * connection.
 	 *
 	 * @author James Moger
 	 *
 	 */
 	static class FanoutNioConnection extends FanoutServiceConnection {
 		final ByteBuffer readBuffer;
 		final ByteBuffer writeBuffer;
 		final List<String> requestQueue;
 		final List<String> replyQueue;
 		final CharsetDecoder decoder;

 		FanoutNioConnection(SocketChannel ch) {
 			super(ch.socket());
 			readBuffer = ByteBuffer.allocate(FanoutConstants.BUFFER_LENGTH);
 			writeBuffer = ByteBuffer.allocate(FanoutConstants.BUFFER_LENGTH);
 			requestQueue = new ArrayList<String>();
 			replyQueue = new ArrayList<String>();
 			decoder = Charset.forName(FanoutConstants.CHARSET).newDecoder();
 		}

 		protected void read(SocketChannel ch, boolean strictRequestTermination) throws CharacterCodingException, IOException {
 			long bytesRead = 0;
 			readBuffer.clear();
 			bytesRead = ch.read(readBuffer);
 			readBuffer.flip();
 			if (bytesRead == -1) {
 				throw new IOException("lost client connection, end of stream");
 			}
 			if (readBuffer.limit() == 0) {
 				return;
 			}
 			CharBuffer cbuf = decoder.decode(readBuffer);
 			String req = cbuf.toString();
 			String [] lines = req.split(strictRequestTermination ? "\n" : "\n|\r");
 			requestQueue.addAll(Arrays.asList(lines));
 		}

 		protected void write(SocketChannel ch) throws IOException {
 			Iterator<String> itr = replyQueue.iterator();
 			while (itr.hasNext()) {
 				String reply = itr.next();
 				writeBuffer.clear();
 				logger.debug(MessageFormat.format("fanout reply to {0}: {1}", id, reply));
 				byte [] bytes = reply.getBytes(FanoutConstants.CHARSET);
 				writeBuffer.put(bytes);
 				if (bytes[bytes.length - 1] != 0xa) {
 					writeBuffer.put((byte) 0xa);
 				}
 				writeBuffer.flip();

 				// loop until write buffer has been completely sent
 				int written = 0;
 				int toWrite = writeBuffer.remaining();
 				while (written != toWrite) {
 					written += ch.write(writeBuffer);
 					try {
 						Thread.sleep(10);
 					} catch (Exception x) {
 					}
 				}
 				itr.remove();
 			}
 			writeBuffer.clear();
 		}

 		@Override
 		protected void reply(String content) throws IOException {
 			// queue the reply
 			// replies are transmitted asynchronously from the requests
 			replyQueue.add(content);
 		}
 	}
 }
	/*
	* Copyright 2013 gitblit.com.
	*
	* Licensed 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.
	*/
	package com.gitblit.fanout;

	import java.io.IOException;
	import java.net.InetSocketAddress;
	import java.nio.ByteBuffer;
	import java.nio.CharBuffer;
	import java.nio.channels.SelectionKey;
	import java.nio.channels.Selector;
	import java.nio.channels.ServerSocketChannel;
	import java.nio.channels.SocketChannel;
	import java.nio.charset.CharacterCodingException;
	import java.nio.charset.Charset;
	import java.nio.charset.CharsetDecoder;
	import java.text.MessageFormat;
	import java.util.ArrayList;
	import java.util.Arrays;
	import java.util.Collection;
	import java.util.HashMap;
	import java.util.Iterator;
	import java.util.List;
	import java.util.Map;
	import java.util.Set;

	import org.slf4j.Logger;
	import org.slf4j.LoggerFactory;

	/**
	* A single-thread NIO implementation of https://github.com/travisghansen/fanout
	*
	* This implementation uses channels and selectors, which are the Java analog of
	* the Linux epoll mechanism used in the original fanout C code.
	*
	* @author James Moger
	*
	*/
	public class FanoutNioService extends FanoutService {

	private final static Logger logger = LoggerFactory.getLogger(FanoutNioService.class);

	private volatile ServerSocketChannel serviceCh;
	private volatile Selector selector;

	public static void main(String[] args) throws Exception {
	FanoutNioService pubsub = new FanoutNioService(null, DEFAULT_PORT);
	pubsub.setStrictRequestTermination(false);
	pubsub.setAllowAllChannelAnnouncements(false);
	pubsub.start();
	}

	/**
	* Create a single-threaded fanout service.
	*
	* @param host
	* @param port
	* the port for running the fanout PubSub service
	* @throws IOException
	*/
	public FanoutNioService(int port) {
	this(null, port);
	}

	/**
	* Create a single-threaded fanout service.
	*
	* @param bindInterface
	* the ip address to bind for the service, may be null
	* @param port
	* the port for running the fanout PubSub service
	* @throws IOException
	*/
	public FanoutNioService(String bindInterface, int port) {
	super(bindInterface, port, "Fanout nio service");
	}

	@Override
	protected boolean isConnected() {
	return serviceCh != null;
	}

	@Override
	protected boolean connect() {
	if (serviceCh == null) {
	try {
	serviceCh = ServerSocketChannel.open();
	serviceCh.configureBlocking(false);
	serviceCh.socket().setReuseAddress(true);
	serviceCh.socket().bind(host == null ? new InetSocketAddress(port) : new InetSocketAddress(host, port));
	selector = Selector.open();
	serviceCh.register(selector, SelectionKey.OP_ACCEPT);
	logger.info(MessageFormat.format("{0} is ready on {1}:{2,number,0}",
	name, host == null ? "0.0.0.0" : host, port));
	} catch (IOException e) {
	logger.error(MessageFormat.format("failed to open {0} on {1}:{2,number,0}",
	name, name, host == null ? "0.0.0.0" : host, port), e);
	return false;
	}
	}
	return true;
	}

	@Override
	protected void disconnect() {
	try {
	if (serviceCh != null) {
	// close all active client connections
	Map<String, SocketChannel> clients = getCurrentClientSockets();
	for (Map.Entry<String, SocketChannel> client : clients.entrySet()) {
	closeClientSocket(client.getKey(), client.getValue());
	}

	// close service socket channel
	logger.debug(MessageFormat.format("closing {0} socket channel", name));
	serviceCh.socket().close();
	serviceCh.close();
	serviceCh = null;
	selector.close();
	selector = null;
	}
	} catch (IOException e) {
	logger.error(MessageFormat.format("failed to disconnect {0}", name), e);
	}
	}

	@Override
	protected void listen() throws IOException {
	while (selector.select(serviceTimeout) > 0) {
	Set<SelectionKey> keys = selector.selectedKeys();
	Iterator<SelectionKey> keyItr = keys.iterator();
	while (keyItr.hasNext()) {
	SelectionKey key = (SelectionKey) keyItr.next();
	if (key.isAcceptable()) {
	// new fanout client connection
	ServerSocketChannel sch = (ServerSocketChannel) key.channel();
	try {
	SocketChannel ch = sch.accept();
	ch.configureBlocking(false);
	configureClientSocket(ch.socket());

	FanoutNioConnection connection = new FanoutNioConnection(ch);
	addConnection(connection);

	// register to send the queued message
	ch.register(selector, SelectionKey.OP_WRITE, connection);
	} catch (IOException e) {
	logger.error("error accepting fanout connection", e);
	}
	} else if (key.isReadable()) {
	// read fanout client request
	SocketChannel ch = (SocketChannel) key.channel();
	FanoutNioConnection connection = (FanoutNioConnection) key.attachment();
	try {
	connection.read(ch, isStrictRequestTermination());
	int replies = 0;
	Iterator<String> reqItr = connection.requestQueue.iterator();
	while (reqItr.hasNext()) {
	String req = reqItr.next();
	String reply = processRequest(connection, req);
	reqItr.remove();
	if (reply != null) {
	replies++;
	}
	}

	if (replies > 0) {
	// register to send the replies to requests
	ch.register(selector, SelectionKey.OP_WRITE, connection);
	} else {
	// re-register for next read
	ch.register(selector, SelectionKey.OP_READ, connection);
	}
	} catch (IOException e) {
	logger.error(MessageFormat.format("fanout connection {0} error: {1}", connection.id, e.getMessage()));
	removeConnection(connection);
	closeClientSocket(connection.id, ch);
	}
	} else if (key.isWritable()) {
	// asynchronous reply to fanout client request
	SocketChannel ch = (SocketChannel) key.channel();
	FanoutNioConnection connection = (FanoutNioConnection) key.attachment();
	try {
	connection.write(ch);

	if (hasConnection(connection)) {
	// register for next read
	ch.register(selector, SelectionKey.OP_READ, connection);
	} else {
	// Connection was rejected due to load or
	// some other reason. Close it.
	closeClientSocket(connection.id, ch);
	}
	} catch (IOException e) {
	logger.error(MessageFormat.format("fanout connection {0}: {1}", connection.id, e.getMessage()));
	removeConnection(connection);
	closeClientSocket(connection.id, ch);
	}
	}
	keyItr.remove();
	}
	}
	}

	protected void closeClientSocket(String id, SocketChannel ch) {
	try {
	ch.close();
	} catch (IOException e) {
	logger.error(MessageFormat.format("fanout connection {0}", id), e);
	}
	}

	protected void broadcast(Collection<FanoutServiceConnection> connections, String channel, String message) {
	super.broadcast(connections, channel, message);

	// register queued write
	Map<String, SocketChannel> sockets = getCurrentClientSockets();
	for (FanoutServiceConnection connection : connections) {
	SocketChannel ch = sockets.get(connection.id);
	if (ch == null) {
	logger.warn(MessageFormat.format("fanout connection {0} has been disconnected", connection.id));
	removeConnection(connection);
	continue;
	}
	try {
	ch.register(selector, SelectionKey.OP_WRITE, connection);
	} catch (IOException e) {
	logger.error(MessageFormat.format("failed to register write op for fanout connection {0}", connection.id));
	}
	}
	}

	protected Map<String, SocketChannel> getCurrentClientSockets() {
	Map<String, SocketChannel> sockets = new HashMap<String, SocketChannel>();
	for (SelectionKey key : selector.keys()) {
	if (key.channel() instanceof SocketChannel) {
	SocketChannel ch = (SocketChannel) key.channel();
	String id = FanoutConstants.getRemoteSocketId(ch.socket());
	sockets.put(id, ch);
	}
	}
	return sockets;
	}

	/**
	* FanoutNioConnection handles reading/writing messages from a remote fanout
	* connection.
	*
	* @author James Moger
	*
	*/
	static class FanoutNioConnection extends FanoutServiceConnection {
	final ByteBuffer readBuffer;
	final ByteBuffer writeBuffer;
	final List<String> requestQueue;
	final List<String> replyQueue;
	final CharsetDecoder decoder;

	FanoutNioConnection(SocketChannel ch) {
	super(ch.socket());
	readBuffer = ByteBuffer.allocate(FanoutConstants.BUFFER_LENGTH);
	writeBuffer = ByteBuffer.allocate(FanoutConstants.BUFFER_LENGTH);
	requestQueue = new ArrayList<String>();
	replyQueue = new ArrayList<String>();
	decoder = Charset.forName(FanoutConstants.CHARSET).newDecoder();
	}

	protected void read(SocketChannel ch, boolean strictRequestTermination) throws CharacterCodingException, IOException {
	long bytesRead = 0;
	readBuffer.clear();
	bytesRead = ch.read(readBuffer);
	readBuffer.flip();
	if (bytesRead == -1) {
	throw new IOException("lost client connection, end of stream");
	}
	if (readBuffer.limit() == 0) {
	return;
	}
	CharBuffer cbuf = decoder.decode(readBuffer);
	String req = cbuf.toString();
	String [] lines = req.split(strictRequestTermination ? "\n" : "\n\|\r");
	requestQueue.addAll(Arrays.asList(lines));
	}

	protected void write(SocketChannel ch) throws IOException {
	Iterator<String> itr = replyQueue.iterator();
	while (itr.hasNext()) {
	String reply = itr.next();
	writeBuffer.clear();
	logger.debug(MessageFormat.format("fanout reply to {0}: {1}", id, reply));
	byte [] bytes = reply.getBytes(FanoutConstants.CHARSET);
	writeBuffer.put(bytes);
	if (bytes[bytes.length - 1] != 0xa) {
	writeBuffer.put((byte) 0xa);
	}
	writeBuffer.flip();

	// loop until write buffer has been completely sent
	int written = 0;
	int toWrite = writeBuffer.remaining();
	while (written != toWrite) {
	written += ch.write(writeBuffer);
	try {
	Thread.sleep(10);
	} catch (Exception x) {
	}
	}
	itr.remove();
	}
	writeBuffer.clear();
	}

	@Override
	protected void reply(String content) throws IOException {
	// queue the reply
	// replies are transmitted asynchronously from the requests
	replyQueue.add(content);
	}
	}
	}