/*
Netwrapper - A library for easy networking in Java
Copyright (C) 2014, University of Lugano
This file is part of Netwrapper.
Netwrapper is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License as published by the Free Software Foundation; either
version 2.1 of the License, or (at your option) any later version.
This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the GNU
Lesser General Public License for more details.
You should have received a copy of the GNU Lesser General Public
License along with this library; if not, write to the Free Software
Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA
*/
package ch.usi.dslab.bezerra.netwrapper.tcp;
import java.io.IOException;
import java.nio.ByteBuffer;
import java.nio.channels.CancelledKeyException;
import java.nio.channels.ClosedChannelException;
import java.nio.channels.SelectionKey;
import java.nio.channels.Selector;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import org.apache.logging.log4j.LogManager;
import org.apache.logging.log4j.Logger;
import ch.usi.dslab.bezerra.netwrapper.Message;
public class TCPSender implements Runnable {
Logger logger = LogManager.getLogger(TCPSender.class.getName());
Selector writeSelector;
Thread TCPSenderThread = null;
boolean running = true;
Map<TCPDestination, TCPConnection> connections = new ConcurrentHashMap<TCPDestination, TCPConnection>();
// write selector data
Queue<TCPMessage> pendingSendMessages = new ConcurrentLinkedQueue<TCPMessage>();
public TCPSender() {
try {
writeSelector = Selector.open();
TCPSenderThread = new Thread(this, "TCPSender");
TCPSenderThread.start();
} catch (IOException e) {
e.printStackTrace();
System.exit(1);
}
}
public TCPConnection connect(TCPDestination destinationReference) throws IOException {
TCPConnection newConnection = new TCPConnection(destinationReference.getAddress(), destinationReference.getPort());
connections.put(destinationReference, newConnection);
return newConnection;
}
public void send(Message message, TCPDestination destinationReference) {
try {
TCPConnection connection = connections.get(destinationReference);
if (connection == null) {
connection = connect(destinationReference);
connections.put(destinationReference, connection);
}
send(message, connection);
}
catch (IOException e) {
e.printStackTrace();
System.exit(1);
}
}
public void send(final Message message, final TCPConnection connection) {
pendingSendMessages.add(new TCPMessage(connection, message.getByteBufferWithLengthHeader()));
writeSelector.wakeup();
}
public void multiDestinationSend(Message message, List<TCPDestination> destinations) {
try {
List<TCPConnection> destinationsConnections = new ArrayList<TCPConnection>();
for (TCPDestination destination : destinations) {
TCPConnection destinationConnection = connections.get(destination);
if (destinationConnection == null) {
destinationConnection = connect(destination);
connections.put(destination, destinationConnection);
}
destinationsConnections.add(destinationConnection);
}
multiConnectionSend(message, destinationsConnections);
}
catch (IOException e) {
e.printStackTrace();
System.exit(1);
}
}
public void multiConnectionSend(final Message message, final List<TCPConnection> connections) {
ByteBuffer bb = message.getByteBufferWithLengthHeader();
for (TCPConnection connection : connections)
pendingSendMessages.add(new TCPMessage(connection, bb.duplicate()));
writeSelector.wakeup();
}
public void stop() {
running = false;
}
@Override
public void run() {
try {
while (running) {
int readyKeys;
if (pendingSendMessages.isEmpty())
readyKeys = writeSelector.select(1000);
else
readyKeys = writeSelector.selectNow();
// checking if there is something to write waiting in line, while
// avoiding starvation by
// registering at most one pending channel to write at a time
TCPMessage pendingMessage = pendingSendMessages.poll();
if (pendingMessage != null) {
pendingMessage.connection.sendQueue.add(pendingMessage.serializedContents);
enableWrite(pendingMessage.connection);
}
if (readyKeys == 0)
continue;
Iterator<SelectionKey> i = writeSelector.selectedKeys().iterator();
while (i.hasNext()) {
SelectionKey key = i.next();
i.remove();
if (!key.isValid())
continue;
try {
if (key.isWritable())
processWrite(key);
}
catch (CancelledKeyException e) {
logger.debug("Key cancelled - probably disconnected... Closing channel.");
key.channel().close();
}
}
}
logger.debug("Exiting tcpsender's select loop");
}
catch(IOException e) {
e.printStackTrace();
System.exit(1);
}
}
private void processWrite(SelectionKey key) {
TCPConnection connection = (TCPConnection) key.attachment();
try {
while(connection.sendQueue.isEmpty() == false) {
ByteBuffer messageToSend = connection.sendQueue.peek();
connection.channel.write(messageToSend);
if (messageToSend.remaining() == 0)
connection.sendQueue.remove();
else
break;
}
}
catch(IOException e) {
e.printStackTrace();
System.exit(1);
}
if (connection.sendQueue.isEmpty())
disableWrite(connection);
}
private void enableWrite(TCPConnection connection) {
try {
SelectionKey writeKey = connection.writeSelectionKey;
if (writeKey == null) {
writeKey = connection.channel.register(writeSelector, SelectionKey.OP_WRITE, connection);
connection.writeSelectionKey = writeKey;
}
else {
final int keyInterest = writeKey.interestOps();
writeKey.interestOps(keyInterest | SelectionKey.OP_WRITE);
}
} catch (ClosedChannelException e) {
e.printStackTrace();
System.exit(1);
}
}
private void disableWrite(TCPConnection connection) {
final SelectionKey writeKey = connection.writeSelectionKey;
final int keyInterest = writeKey.interestOps();
writeKey.interestOps(keyInterest & ~SelectionKey.OP_WRITE);
}
}