NioReplicationTask.java
/*
* 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.
*/
package org.apache.catalina.tribes.transport.nio;
import java.io.IOException;
import java.net.SocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.DatagramChannel;
import java.nio.channels.ReadableByteChannel;
import java.nio.channels.SelectionKey;
import java.nio.channels.SocketChannel;
import java.nio.channels.WritableByteChannel;
import org.apache.catalina.tribes.ChannelMessage;
import org.apache.catalina.tribes.ChannelReceiver;
import org.apache.catalina.tribes.RemoteProcessException;
import org.apache.catalina.tribes.UniqueId;
import org.apache.catalina.tribes.io.BufferPool;
import org.apache.catalina.tribes.io.ChannelData;
import org.apache.catalina.tribes.io.ListenCallback;
import org.apache.catalina.tribes.io.ObjectReader;
import org.apache.catalina.tribes.transport.AbstractRxTask;
import org.apache.catalina.tribes.transport.Constants;
import org.apache.catalina.tribes.util.Logs;
import org.apache.catalina.tribes.util.StringManager;
import org.apache.juli.logging.Log;
import org.apache.juli.logging.LogFactory;
/**
* A worker thread class which can drain channels and echo-back the input. Each instance is constructed with a reference
* to the owning thread pool object. When started, the thread loops forever waiting to be awakened to service the
* channel associated with a SelectionKey object. The worker is tasked by calling its serviceChannel() method with a
* SelectionKey object. The serviceChannel() method stores the key reference in the thread object then calls notify() to
* wake it up. When the channel has been drained, the worker thread returns itself to its parent pool.
*/
public class NioReplicationTask extends AbstractRxTask {
private static final Log log = LogFactory.getLog(NioReplicationTask.class);
protected static final StringManager sm = StringManager.getManager(NioReplicationTask.class);
private ByteBuffer buffer = null;
private SelectionKey key;
private int rxBufSize;
private final NioReceiver receiver;
public NioReplicationTask(ListenCallback callback, NioReceiver receiver) {
super(callback);
this.receiver = receiver;
}
// loop forever waiting for work to do
@Override
public synchronized void run() {
if (buffer == null) {
int size = getRxBufSize();
if (key.channel() instanceof DatagramChannel) {
size = ChannelReceiver.MAX_UDP_SIZE;
}
if ((getOptions() & OPTION_DIRECT_BUFFER) == OPTION_DIRECT_BUFFER) {
buffer = ByteBuffer.allocateDirect(size);
} else {
buffer = ByteBuffer.allocate(size);
}
} else {
buffer.clear();
}
if (key == null) {
return; // just in case
}
if (log.isTraceEnabled()) {
log.trace("Servicing key:" + key);
}
try {
ObjectReader reader = (ObjectReader) key.attachment();
if (reader == null) {
if (log.isTraceEnabled()) {
log.trace("No object reader, cancelling:" + key);
}
cancelKey(key);
} else {
if (log.isTraceEnabled()) {
log.trace("Draining channel:" + key);
}
drainChannel(key, reader);
}
} catch (Exception e) {
// this is common, since the sockets on the other
// end expire after a certain time.
if (e instanceof CancelledKeyException) {
// do nothing
} else if (e instanceof IOException) {
// don't spew out stack traces for IO exceptions unless debug is enabled.
if (log.isDebugEnabled()) {
log.debug(sm.getString("nioReplicationTask.unable.drainChannel.ioe", e.getMessage()), e);
} else {
log.warn(sm.getString("nioReplicationTask.unable.drainChannel.ioe", e.getMessage()));
}
} else if (log.isErrorEnabled()) {
// this is a real error, log it.
log.error(sm.getString("nioReplicationTask.exception.drainChannel"), e);
}
cancelKey(key);
}
key = null;
// done, ready for more, return to pool
getTaskPool().returnWorker(this);
}
/**
* Called to initiate a unit of work by this worker thread on the provided SelectionKey object. This method is
* synchronized, as is the run() method, so only one key can be serviced at a given time. Before waking the worker
* thread, and before returning to the main selection loop, this key's interest set is updated to remove OP_READ.
* This will cause the selector to ignore read-readiness for this channel while the worker thread is servicing it.
*
* @param key The key to process
*/
public synchronized void serviceChannel(SelectionKey key) {
if (log.isTraceEnabled()) {
log.trace("About to service key:" + key);
}
ObjectReader reader = (ObjectReader) key.attachment();
if (reader != null) {
reader.setLastAccess(System.currentTimeMillis());
}
this.key = key;
key.interestOps(key.interestOps() & (~SelectionKey.OP_READ));
key.interestOps(key.interestOps() & (~SelectionKey.OP_WRITE));
}
/**
* The actual code which drains the channel associated with the given key. This method assumes the key has been
* modified prior to invocation to turn off selection interest in OP_READ. When this method completes it re-enables
* OP_READ and calls wakeup() on the selector so the selector will resume watching this channel.
*
* @param key The key to process
* @param reader The reader
*
* @throws Exception IO error
*/
protected void drainChannel(final SelectionKey key, ObjectReader reader) throws Exception {
reader.access();
ReadableByteChannel channel = (ReadableByteChannel) key.channel();
int count = -1;
SocketAddress saddr = null;
if (channel instanceof SocketChannel) {
// loop while data available, channel is non-blocking
while ((count = channel.read(buffer)) > 0) {
buffer.flip(); // make buffer readable
if (buffer.hasArray()) {
reader.append(buffer.array(), 0, count, false);
} else {
reader.append(buffer, count, false);
}
buffer.clear(); // make buffer empty
// do we have at least one package?
if (reader.hasPackage()) {
break;
}
}
} else if (channel instanceof DatagramChannel) {
DatagramChannel dchannel = (DatagramChannel) channel;
saddr = dchannel.receive(buffer);
buffer.flip(); // make buffer readable
if (buffer.hasArray()) {
reader.append(buffer.array(), 0, buffer.limit() - buffer.position(), false);
} else {
reader.append(buffer, buffer.limit() - buffer.position(), false);
}
buffer.clear(); // make buffer empty
// did we get a package
count = reader.hasPackage() ? 1 : -1;
}
int pkgcnt = reader.count();
if (count < 0 && pkgcnt == 0) {
// end of stream, and no more packages to process
remoteEof(key);
return;
}
ChannelMessage[] msgs = pkgcnt == 0 ? ChannelData.EMPTY_DATA_ARRAY : reader.execute();
registerForRead(key, reader);// register to read new data, before we send it off to avoid dead locks
for (ChannelMessage msg : msgs) {
/*
* Use send ack here if you want to ack the request to the remote server before completing the request This
* is considered an asynchronous request
*/
if (ChannelData.sendAckAsync(msg.getOptions())) {
sendAck(key, (WritableByteChannel) channel, Constants.ACK_COMMAND, saddr);
}
try {
if (Logs.MESSAGES.isTraceEnabled()) {
try {
Logs.MESSAGES.trace("NioReplicationThread - Received msg:" + new UniqueId(msg.getUniqueId()) +
" at " + new java.sql.Timestamp(System.currentTimeMillis()));
} catch (Throwable t) {
}
}
// process the message
getCallback().messageDataReceived(msg);
/*
* Use send ack here if you want the request to complete on this server before sending the ack to the
* remote server This is considered a synchronized request
*/
if (ChannelData.sendAckSync(msg.getOptions())) {
sendAck(key, (WritableByteChannel) channel, Constants.ACK_COMMAND, saddr);
}
} catch (RemoteProcessException e) {
if (log.isDebugEnabled()) {
log.error(sm.getString("nioReplicationTask.process.clusterMsg.failed"), e);
}
if (ChannelData.sendAckSync(msg.getOptions())) {
sendAck(key, (WritableByteChannel) channel, Constants.FAIL_ACK_COMMAND, saddr);
}
} catch (Exception e) {
log.error(sm.getString("nioReplicationTask.process.clusterMsg.failed"), e);
if (ChannelData.sendAckSync(msg.getOptions())) {
sendAck(key, (WritableByteChannel) channel, Constants.FAIL_ACK_COMMAND, saddr);
}
}
if (getUseBufferPool()) {
BufferPool.getBufferPool().returnBuffer(msg.getMessage());
msg.setMessage(null);
}
}
if (count < 0) {
remoteEof(key);
}
}
private void remoteEof(SelectionKey key) {
// close channel on EOF, invalidates the key
if (log.isDebugEnabled()) {
log.debug(sm.getString("nioReplicationTask.disconnect"));
}
cancelKey(key);
}
protected void registerForRead(final SelectionKey key, ObjectReader reader) {
if (log.isTraceEnabled()) {
log.trace("Adding key for read event:" + key);
}
reader.finish();
// register our OP_READ interest
Runnable r = () -> {
try {
if (key.isValid()) {
// resume interest in OP_READ, OP_WRITE
int resumeOps = key.interestOps() | SelectionKey.OP_READ;
key.interestOps(resumeOps);
if (log.isTraceEnabled()) {
log.trace("Registering key for read:" + key);
}
}
} catch (CancelledKeyException ckx) {
NioReceiver.cancelledKey(key);
if (log.isTraceEnabled()) {
log.trace("CKX Cancelling key:" + key);
}
} catch (Exception x) {
log.error(sm.getString("nioReplicationTask.error.register.key", key), x);
}
};
receiver.addEvent(r);
}
private void cancelKey(final SelectionKey key) {
if (log.isTraceEnabled()) {
log.trace("Adding key for cancel event:" + key);
}
ObjectReader reader = (ObjectReader) key.attachment();
if (reader != null) {
reader.setCancelled(true);
reader.finish();
}
Runnable cx = () -> {
if (log.isTraceEnabled()) {
log.trace("Cancelling key:" + key);
}
NioReceiver.cancelledKey(key);
};
receiver.addEvent(cx);
}
/**
* Send a reply-acknowledgement (6,2,3), sends it doing a busy write, the ACK is so small that it should always go
* to the buffer.
*
* @param key The key to use
* @param channel The channel
* @param command The command to write
* @param udpaddr Target address
*/
protected void sendAck(SelectionKey key, WritableByteChannel channel, byte[] command, SocketAddress udpaddr) {
try {
ByteBuffer buf = ByteBuffer.wrap(command);
int total = 0;
if (channel instanceof DatagramChannel) {
DatagramChannel dchannel = (DatagramChannel) channel;
// were using a shared channel, document says its thread safe
// TODO check optimization, one channel per thread?
while (total < command.length) {
total += dchannel.send(buf, udpaddr);
}
} else {
while (total < command.length) {
total += channel.write(buf);
}
}
if (log.isTraceEnabled()) {
log.trace("ACK sent to " +
((channel instanceof SocketChannel) ? ((SocketChannel) channel).socket().getInetAddress() :
((DatagramChannel) channel).socket().getInetAddress()));
}
} catch (IOException x) {
log.warn(sm.getString("nioReplicationTask.unable.ack", x.getMessage()));
}
}
public void setRxBufSize(int rxBufSize) {
this.rxBufSize = rxBufSize;
}
public int getRxBufSize() {
return rxBufSize;
}
}