public static void chooseDispatcher(SocketChannel channel, int op) {
if (channel.isRegistered()) {
for (Dispatcher dispatcher : dispatchers) {
SelectionKey key = channel.keyFor(dispatcher.getSelector());
if (key != null && key.isValid() && (key.interestOps() == op)) {
return;
}
}
for (Dispatcher dispatcher : dispatchers) {
SelectionKey key = channel.keyFor(dispatcher.getSelector());
if (key != null && key.isValid() && (key.interestOps() == 0)) {
if (op == SelectionKey.OP_READ) {
dispatcher.addChannelForRead(channel);
} else if (op == SelectionKey.OP_WRITE) {
dispatcher.addChannelForWrite(channel);
}
return;
}
}
} else {
if (op == SelectionKey.OP_READ) {
nextDispatcher().addChannelForRead(channel);
} else if (op == SelectionKey.OP_WRITE) {
nextDispatcher().addChannelForWrite(channel);
}
}
}
public static void main(String[] args) throws Exception {
SocketChannel sc = SocketChannel.open();
Selector sel = Selector.open();
if (sc.keyFor(sel) != null) throw new Exception("keyFor != null");
sc.configureBlocking(false);
SelectionKey sk = sc.register(sel, SelectionKey.OP_READ, args);
if (sc.keyFor(sel) != sk) throw new Exception("keyFor returned " + sc.keyFor(sel));
if (sk.attachment() != args) throw new Exception("attachment() returned " + sk.attachment());
Trivial t = new Trivial();
sk.attach(t);
if (sk.attachment() != t) throw new Exception("Wrong attachment");
sk.isReadable();
sk.isWritable();
sk.isConnectable();
sk.isAcceptable();
}
@Override
public void finishedReading(SelectionKey key) {
((AttachmentHTTP) serverChannel.keyFor(key.selector()).attachment()).fullWrittenBuffer = true;
PoolClientConnectionManager.getInstance().removeClient(key);
Statistics.getInstance().addMethod(httpRequest.getMethod());
if (httpRequest.getHost() != null) {
LoggingManager.logReport("REQ - " + httpRequest.getMethod() + " " + httpRequest.getHost());
} else {
LoggingManager.logError("REQ - " + httpRequest.getMethod() + " " + "Unkown host");
}
}
protected void cleanUp(SocketChannel socketChannel) {
try {
socketChannel.close();
} catch (IOException e) {
logger.info(Util.getErrorMessage(e));
}
SelectionKey key = socketChannel.keyFor(_selector);
if (key != null) {
key.cancel();
}
if (_pendingData.containsKey(socketChannel)) {
_pendingData.remove(socketChannel);
}
}
protected void clearOpWrite(SocketChannel channel) {
SelectionKey key = channel.keyFor(nioSelector);
if (key == null) {
return;
}
if (!key.isValid()) {
close(key);
return;
}
int interestOps = key.interestOps();
if ((interestOps & SelectionKey.OP_WRITE) != 0) {
interestOps &= ~SelectionKey.OP_WRITE;
key.interestOps(interestOps);
}
}
public void run() {
try {
if (log.logTrace()) {
log.trace("Idle client timer expired: " + System.identityHashCode(socket));
}
SocketChannel socketChannel = (SocketChannel) this.socket.getChannel();
socketChannel.keyFor(HttpRpcServer.this.getSocketSelector()).cancel();
// This (shutting down the output stream) seems unnecessary but
// without it the client never sees a disconnect under Linux.
// For good measure we shutdown the input stream too.
socketChannel.socket().shutdownOutput();
socketChannel.socket().shutdownInput();
socketChannel.close();
HttpRpcServer.this.deregisterSocket(socket);
} catch (Exception e) {
log.warn("IdleClientTimerTask caught an exception", e);
}
}
/**
* removes back in the OP_WRITE from the selector, otherwise it'll spin loop. The OP_WRITE will
* get added back in as soon as we have data to write
*
* @param socketChannel the socketChannel we wish to stop writing to
* @param attached data associated with the socketChannels key
*/
public synchronized void disableWrite(
@NotNull final SocketChannel socketChannel, @NotNull final Attached attached) {
try {
SelectionKey key = socketChannel.keyFor(selector);
if (key != null) {
if (attached.isHandShakingComplete() && selector.isOpen()) {
if (LOG.isDebugEnabled())
LOG.debug(
"Disabling OP_WRITE to remoteIdentifier="
+ attached.remoteIdentifier
+ ", localIdentifier="
+ localIdentifier);
key.interestOps(key.interestOps() & ~OP_WRITE);
}
}
} catch (Exception e) {
LOG.error("", e);
}
}
public void close(SocketChannel ch) {
if (!listener.handleConnectionClose(ch)) return;
try {
ch.close();
} catch (IOException e) {
e.printStackTrace();
}
ch.keyFor(selector).cancel();
synchronized (changeRequests) {
Iterator changes = changeRequests.iterator();
while (changes.hasNext()) {
ChangeRequest req = (ChangeRequest) changes.next();
if (req.socket == ch) {
changeRequests.remove(req);
break;
}
}
}
}
public boolean hasChannel(SocketChannel ch) {
return ch != null && ch.keyFor(selector) != null;
}
private void closeConnection(SocketChannel socketChannel) throws IOException {
socketChannel.close();
socketChannel.keyFor(selector).cancel();
}
private void readBinary(Connection c) throws IOException {
sc = c.channel();
buffer.clear();
oldLimit = 0;
bufferSize = 0;
read = 0;
totalRead = 0;
if (sc == null || c.isWriteDisabled() || !sc.isOpen()) {
c.close(false);
}
try {
synchronized (sc) {
totalRead = read = sc.read(buffer);
}
} catch (IOException e) {
if (debugEnabled) log.debug("IOException reading from connection " + c, e);
c.close(false);
}
if (read > 0) {
bufferSize = buffer.getShort(0);
tries = 0;
if (debugEnabled)
log.debug(
"Connection " + c + " about to process " + buffer + ", packet size: " + bufferSize);
while (buffer.position() < bufferSize && tries < maxTries) {
tries++;
before = System.currentTimeMillis();
synchronized (sc) {
totalRead += read = sc.read(buffer);
}
if (read == 0) {
if (sc.keyFor(retrySelector) == null) {
sc.register(retrySelector, SelectionKey.OP_READ);
} else {
sc.keyFor(retrySelector).interestOps(SelectionKey.OP_READ);
}
retrySelection = retrySelector.select();
if (retrySelection > 0) {
retrySelector.selectedKeys().clear();
}
}
}
if (sc.keyFor(retrySelector) != null) {
sc.keyFor(retrySelector).interestOps(0);
}
if (tries == maxTries) {
if (debugEnabled)
log.error(
"Too much read tries ("
+ maxTries
+ ") without any bytes read (read: "
+ totalRead
+ ", remaining: "
+ buffer.remaining()
+ ") for connection "
+ c
+ ", kicking client");
c.close(false);
return;
} else if (tries > 0 && debugEnabled) {
timeLostInRetries += System.currentTimeMillis() - before;
log.debug(
"Read successfully "
+ totalRead
+ " bytes after "
+ tries
+ " tries (total time lost: "
+ timeLostInRetries
+ " ms)");
}
buffer.flip();
while (buffer.remaining() > 2 && buffer.remaining() >= buffer.getShort(buffer.position())) {
try {
bufferSize = buffer.getShort();
if (bufferSize > 1) {
bufferSize -= 2;
}
oldLimit = buffer.limit() - 2;
buffer.compact();
buffer.limit(bufferSize);
buffer.position(0);
} catch (IllegalArgumentException e) {
if (debugEnabled)
log.debug(
"Illegal argument while parsing buffer "
+ buffer
+ " (read: "
+ totalRead
+ ", awaited: "
+ bufferSize
+ ") for connection "
+ c,
e);
c.close(false);
return;
}
if (debugEnabled)
log.debug("Connection " + c + " about to process data from buffer " + buffer);
if (workersEnabled) {
try {
choose().add(c, buffer);
} catch (RuntimeException e) {
if (debugEnabled)
log.debug("Cannot queue packet " + buffer + " for connection " + c, e);
c.close(false);
}
} else {
try {
if (c.processData(buffer)) {
if (buffer.position() != bufferSize) {
if (debugEnabled)
log.debug(
"After processing, buffer position is not as expected: expected "
+ (bufferSize)
+ ", buffer: "
+ buffer
+ ", fixing...");
buffer.position(bufferSize);
}
} else {
c.close(false);
return;
}
} catch (Exception e) {
if (debugEnabled)
log.error(
e.getClass().getSimpleName()
+ " while processing buffer "
+ buffer
+ " for connection "
+ c,
e);
c.close(false);
}
}
if (oldLimit > buffer.position()) {
buffer.limit(oldLimit);
if (debugEnabled)
log.debug(
"Connection "
+ c
+ ": buffer "
+ buffer
+ " has more packets (old limit: "
+ oldLimit
+ ", last packet size: "
+ bufferSize
+ ", next packet size: "
+ buffer.getShort(buffer.position())
+ ") ...");
buffer.compact();
buffer.position(0);
buffer.limit(oldLimit - bufferSize);
if (debugEnabled)
log.debug(
"Connection "
+ c
+ " about to process next packet from buffer "
+ buffer
+ ", next packet size: "
+ buffer.getShort(0));
} else {
if (debugEnabled)
log.debug(
"Connection "
+ c
+ " buffer "
+ buffer
+ " seems entirely read, old limit: "
+ oldLimit
+ ", read: "
+ totalRead);
break;
}
}
if (buffer.hasRemaining()) {
if (debugEnabled)
log.error(
"Buffer "
+ buffer
+ " still has data (awaited: "
+ bufferSize
+ ", read: "
+ totalRead
+ "), discarding and closing connection "
+ c
+ "...");
c.close(false);
return;
}
} else if (read < 0) {
c.close(false);
}
}
@Override
public void run() {
// loop while any socket is open
while (incoming.isOpen() || !peers.isEmpty()) {
try {
// wait on selector
selector.select(10000);
// Handle all ready channels
Set<SelectionKey> selected = selector.selectedKeys();
for (SelectionKey k : selected.toArray(new SelectionKey[0])) {
selected.remove(k); // We're handling it
if (!k.isValid()) continue; // Invalid?
Channel c = k.channel();
if (c.equals(incoming)) {
acceptIncomingPeer();
} else if (new_sockets.contains(c)) {
processNewSocket((SocketChannel) c);
} else if (new_peers.containsKey(c)) {
processNewPeer((SocketChannel) c);
} else if (peers.containsKey(c)) {
processPeerMessage((SocketChannel) c);
} else {
System.err.println("Tried to process unknown socket.");
k.cancel(); // If we don't know it now, we'll probably never know it.
c.close();
}
}
// Check for closed sockets
for (SocketChannel socket : peers.keySet().toArray(new SocketChannel[0])) {
if (!socket.isOpen()) {
socket.keyFor(selector).cancel();
Peer peer = peers.get(socket);
peers.remove(socket);
sockets.remove(peer);
if (listener != null) listener.dropPeer(peer);
// TODO: Reconnect, propose drop
}
}
for (SocketChannel socket : new_sockets) {
if (!socket.isOpen()) {
socket.keyFor(selector).cancel();
new_sockets.remove(socket);
}
}
for (SocketChannel socket : new_peers.keySet()) {
if (!socket.isOpen()) {
socket.keyFor(selector).cancel();
new_peers.remove(socket);
// TODO: Retry?
} else if (!socket.isRegistered()) {
// Check for new sockets.
socket.register(selector, SelectionKey.OP_READ);
}
}
} catch (IOException e) {
// TODO Handle this better. In mean time, just keep going.
System.err.println(e);
e.printStackTrace();
}
}
try {
selector.close();
} catch (IOException e) {
// Not much to do.
e.printStackTrace();
}
}