net.data.technology.jraft.extensions.RpcTcpClient Maven / Gradle / Ivy
Go to download
Show more of this group Show more artifacts with this name
Show all versions of jraft-exts Show documentation
Show all versions of jraft-exts Show documentation
extension project for Jraft (Java implementation of Raft Consensus), this module contains may default implementations for interfaces that required by Jraft
/**
*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. 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 net.data.technology.jraft.extensions;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.nio.ByteBuffer;
import java.nio.channels.AsynchronousChannelGroup;
import java.nio.channels.AsynchronousSocketChannel;
import java.nio.channels.CompletionHandler;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.ExecutorService;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import org.apache.log4j.LogManager;
import org.apache.log4j.Logger;
import net.data.technology.jraft.RaftRequestMessage;
import net.data.technology.jraft.RaftResponseMessage;
import net.data.technology.jraft.RpcClient;
public class RpcTcpClient implements RpcClient {
private AsynchronousSocketChannel connection;
private AsynchronousChannelGroup channelGroup;
private ConcurrentLinkedQueue> readTasks;
private ConcurrentLinkedQueue> writeTasks;
private AtomicInteger readers;
private AtomicInteger writers;
private InetSocketAddress remote;
private Logger logger;
public RpcTcpClient(InetSocketAddress remote, ExecutorService executorService){
this.remote = remote;
this.logger = LogManager.getLogger(getClass());
this.readTasks = new ConcurrentLinkedQueue>();
this.writeTasks = new ConcurrentLinkedQueue>();
this.readers = new AtomicInteger(0);
this.writers = new AtomicInteger(0);
try{
this.channelGroup = AsynchronousChannelGroup.withThreadPool(executorService);
}catch(IOException err){
this.logger.error("failed to create channel group", err);
throw new RuntimeException("failed to create the channel group due to errors.");
}
}
@Override
public synchronized CompletableFuture send(final RaftRequestMessage request) {
this.logger.debug(String.format("trying to send message %s to server %d at endpoint %s", request.getMessageType().toString(), request.getDestination(), this.remote.toString()));
CompletableFuture result = new CompletableFuture();
if(this.connection == null || !this.connection.isOpen()){
try{
this.connection = AsynchronousSocketChannel.open(this.channelGroup);
this.connection.connect(this.remote, new AsyncTask(request, result), handlerFrom((Void v, AsyncTask task) -> {
sendAndRead(task, false);
}));
}catch(Throwable error){
closeSocket();
result.completeExceptionally(error);
}
}else{
this.sendAndRead(new AsyncTask(request, result), false);
}
return result;
}
private void sendAndRead(AsyncTask task, boolean skipQueueing){
if(!skipQueueing){
int writerCount = this.writers.getAndIncrement();
if(writerCount > 0){
this.logger.debug("there is a pending write, queue this write task");
this.writeTasks.add(task);
return;
}
}
ByteBuffer buffer = ByteBuffer.wrap(BinaryUtils.messageToBytes(task.input));
try{
AsyncUtility.writeToChannel(this.connection, buffer, task, handlerFrom((Integer bytesSent, AsyncTask context) -> {
if(bytesSent.intValue() < buffer.limit()){
logger.info("failed to sent the request to remote server.");
context.future.completeExceptionally(new IOException("Only partial of the data could be sent"));
closeSocket();
}else{
// read the response
ByteBuffer responseBuffer = ByteBuffer.allocate(BinaryUtils.RAFT_RESPONSE_HEADER_SIZE);
this.readResponse(new AsyncTask(responseBuffer, context.future), false);
}
int waitingWriters = this.writers.decrementAndGet();
if(waitingWriters > 0){
this.logger.debug("there are pending writers in queue, will try to process them");
AsyncTask pendingTask = null;
while((pendingTask = this.writeTasks.poll()) == null);
this.sendAndRead(pendingTask, true);
}
}));
}catch(Exception writeError){
logger.info("failed to write the socket", writeError);
task.future.completeExceptionally(writeError);
closeSocket();
}
}
private void readResponse(AsyncTask task, boolean skipQueueing){
if(!skipQueueing){
int readerCount = this.readers.getAndIncrement();
if(readerCount > 0){
this.logger.debug("there is a pending read, queue this read task");
this.readTasks.add(task);
return;
}
}
CompletionHandler> handler = handlerFrom((Integer bytesRead, AsyncTask context) -> {
if(bytesRead.intValue() < BinaryUtils.RAFT_RESPONSE_HEADER_SIZE){
logger.info("failed to read response from remote server.");
context.future.completeExceptionally(new IOException("Only part of the response data could be read"));
closeSocket();
}else{
RaftResponseMessage response = BinaryUtils.bytesToResponseMessage(context.input.array());
context.future.complete(response);
}
int waitingReaders = this.readers.decrementAndGet();
if(waitingReaders > 0){
this.logger.debug("there are pending readers in queue, will try to process them");
AsyncTask pendingTask = null;
while((pendingTask = this.readTasks.poll()) == null);
this.readResponse(pendingTask, true);
}
});
try{
this.logger.debug("reading response from socket...");
AsyncUtility.readFromChannel(this.connection, task.input, task, handler);
}catch(Exception readError){
logger.info("failed to read from socket", readError);
task.future.completeExceptionally(readError);
closeSocket();
}
}
private CompletionHandler> handlerFrom(BiConsumer> completed) {
return AsyncUtility.handlerFrom(completed, (Throwable error, AsyncTask context) -> {
this.logger.info("socket error", error);
context.future.completeExceptionally(error);
closeSocket();
});
}
private synchronized void closeSocket(){
this.logger.debug("close the socket due to errors");
try{
if(this.connection != null){
this.connection.close();
this.connection = null;
}
}catch(IOException ex){
this.logger.info("failed to close socket", ex);
}
while(true){
AsyncTask task = this.readTasks.poll();
if(task == null){
break;
}
task.future.completeExceptionally(new IOException("socket is closed, no more reads can be completed"));
}
while(true){
AsyncTask task = this.writeTasks.poll();
if(task == null){
break;
}
task.future.completeExceptionally(new IOException("socket is closed, no more writes can be completed"));
}
this.readers.set(0);
this.writers.set(0);
}
static class AsyncTask{
private TInput input;
private CompletableFuture future;
public AsyncTask(TInput input, CompletableFuture future){
this.input = input;
this.future = future;
}
}
}