use of org.apache.ratis.thirdparty.io.netty.channel.ChannelHandlerContext in project incubator-ratis by apache.
the class DataStreamManagement method readImpl.
private void readImpl(DataStreamRequestByteBuf request, ChannelHandlerContext ctx, ByteBuf buf, CheckedBiFunction<RaftClientRequest, Set<RaftPeer>, Set<DataStreamOutputRpc>, IOException> getStreams) {
boolean close = WriteOption.containsOption(request.getWriteOptions(), StandardWriteOption.CLOSE);
ClientInvocationId key = ClientInvocationId.valueOf(request.getClientId(), request.getStreamId());
final StreamInfo info;
if (request.getType() == Type.STREAM_HEADER) {
final MemoizedSupplier<StreamInfo> supplier = JavaUtils.memoize(() -> newStreamInfo(buf, getStreams));
info = streams.computeIfAbsent(key, id -> supplier.get());
if (!supplier.isInitialized()) {
streams.remove(key);
throw new IllegalStateException("Failed to create a new stream for " + request + " since a stream already exists Key: " + key + " StreamInfo:" + info);
}
getMetrics().onRequestCreate(RequestType.HEADER);
} else if (close) {
info = Optional.ofNullable(streams.remove(key)).orElseThrow(() -> new IllegalStateException("Failed to remove StreamInfo for " + request));
} else {
info = Optional.ofNullable(streams.get(key)).orElseThrow(() -> {
streams.remove(key);
return new IllegalStateException("Failed to get StreamInfo for " + request);
});
}
final CompletableFuture<Long> localWrite;
final List<CompletableFuture<DataStreamReply>> remoteWrites;
if (request.getType() == Type.STREAM_HEADER) {
localWrite = CompletableFuture.completedFuture(0L);
remoteWrites = Collections.emptyList();
} else if (request.getType() == Type.STREAM_DATA) {
localWrite = info.getLocal().write(buf, request.getWriteOptions(), writeExecutor);
remoteWrites = info.applyToRemotes(out -> out.write(request, requestExecutor));
} else {
throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
}
composeAsync(info.getPrevious(), requestExecutor, n -> JavaUtils.allOf(remoteWrites).thenCombineAsync(localWrite, (v, bytesWritten) -> {
if (request.getType() == Type.STREAM_HEADER || (request.getType() == Type.STREAM_DATA && !close)) {
sendReply(remoteWrites, request, bytesWritten, info.getCommitInfos(), ctx);
} else if (close) {
if (info.isPrimary()) {
// after all server close stream, primary server start transaction
startTransaction(info, request, bytesWritten, ctx);
} else {
sendReply(remoteWrites, request, bytesWritten, info.getCommitInfos(), ctx);
}
} else {
throw new IllegalStateException(this + ": Unexpected type " + request.getType() + ", request=" + request);
}
return null;
}, requestExecutor)).whenComplete((v, exception) -> {
try {
if (exception != null) {
streams.remove(key);
replyDataStreamException(server, exception, info.getRequest(), request, ctx);
}
} finally {
buf.release();
}
});
}
Aggregations