org.infinispan.server.resp.RespRequestHandler Maven / Gradle / Ivy
Show all versions of infinispan-server-resp Show documentation
package org.infinispan.server.resp;
import java.util.List;
import java.util.Objects;
import java.util.concurrent.CompletableFuture;
import java.util.concurrent.CompletionStage;
import java.util.concurrent.Executor;
import java.util.function.BiConsumer;
import java.util.function.BiFunction;
import java.util.function.Consumer;
import java.util.function.Function;
import org.infinispan.commons.util.concurrent.CompletionStages;
import org.infinispan.server.resp.commands.connection.QUIT;
import org.infinispan.server.resp.serialization.Resp3Response;
import io.netty.channel.ChannelHandlerContext;
import io.netty.util.AttributeKey;
public abstract class RespRequestHandler {
protected final CompletionStage myStage = CompletableFuture.completedFuture(this);
protected final RespServer respServer;
public static final AttributeKey BYTE_BUF_POOL_ATTRIBUTE_KEY = AttributeKey.newInstance("buffer-pool");
private final BiFunction HANDLE_FAILURE = (h, t) -> {
if (t != null) {
Resp3Response.error(t, allocator());
// If exception, never use handler
return this;
}
return h;
};
ByteBufPool allocatorToUse;
ChannelHandlerContext ctx;
protected RespRequestHandler(RespServer server) {
this.respServer = server;
}
protected void initializeIfNecessary(ChannelHandlerContext ctx) {
if (allocatorToUse == null) {
if (!ctx.channel().hasAttr(BYTE_BUF_POOL_ATTRIBUTE_KEY)) {
throw new IllegalStateException("BufferPool was not initialized in the context " + ctx);
}
allocatorToUse = ctx.channel().attr(BYTE_BUF_POOL_ATTRIBUTE_KEY).get();
this.ctx = ctx;
}
}
public RespServer respServer() {
return respServer;
}
public CompletionStage myStage() {
return myStage;
}
/**
* Acquire the buffer allocator in the current context.
*
* The {@link ByteBufPool} provides the means to allocate a {@link io.netty.buffer.ByteBuf} with the required
* free size. All the writes must utilize the {@link ByteBufPool} instead of manually allocating buffers.
*
* Thread-safety: The allocator and the {@link io.netty.buffer.ByteBuf} must be utilized only from the
* event loop thread. Use the {@link ChannelHandlerContext} to verify.
*
* @return A {@link io.netty.buffer.ByteBuf} allocator.
*/
public ByteBufPool allocator() {
assert ctx.channel().eventLoop().inEventLoop() : "Buffer allocation should occur in event loop, it was " + Thread.currentThread().getName();
return allocatorToUse;
}
public final CompletionStage handleRequest(ChannelHandlerContext ctx, RespCommand command, List arguments) {
initializeIfNecessary(ctx);
if (command == null) {
RespErrorUtil.unknownCommand(allocatorToUse);
return myStage;
}
if (!command.hasValidNumberOfArguments(arguments)) {
RespErrorUtil.wrongArgumentNumber(command, allocator());
return myStage;
}
try {
return actualHandleRequest(ctx, command, arguments);
} catch (Exception e) {
// Extract the handler capable of writing the error message.
Consumer writer = RespErrorUtil.handleException(e);
// If there is no writer capable of handling the message, return the exception to the lower level to handle.
if (writer == null) {
return CompletableFuture.failedFuture(e);
}
writer.accept(allocatorToUse);
return myStage();
}
}
/**
* Handles the RESP request returning a stage that when complete notifies the command has completed as well as
* providing the request handler for subsequent commands.
*
* Implementations should never use the ByteBufAllocator in the context.
* Instead, they should use {@link #allocatorToUse} to retrieve a ByteBuffer.
* This ByteBuffer should only have bytes written to it adding up to the size requested.
* The ByteBuffer itself should never be written to the context or channel and flush should also never be invoked.
* Failure to do so may cause mis-ordering or responses as requests support pipelining and a ByteBuf may not be
* send down stream until later in the pipeline.
*
* @param ctx Netty context pipeline for this request
* @param type The command type
* @param arguments The remaining arguments to the command
* @return stage that when complete returns the new handler to instate. This stage must be completed on the event loop
*/
protected CompletionStage actualHandleRequest(ChannelHandlerContext ctx, RespCommand type, List arguments) {
if (type instanceof QUIT) {
ctx.close();
return myStage;
}
RespErrorUtil.unknownCommand(allocatorToUse);
return myStage;
}
public void handleChannelDisconnect(ChannelHandlerContext ctx) {
// Do nothing by default
}
public CompletionStage stageToReturn(CompletionStage stage, ChannelHandlerContext ctx,
BiConsumer biConsumer) {
return stageToReturn(stage, ctx, Objects.requireNonNull(biConsumer), null);
}
public CompletionStage stageToReturn(CompletionStage stage, ChannelHandlerContext ctx,
Function handlerWhenComplete) {
return stageToReturn(stage, ctx, null, Objects.requireNonNull(handlerWhenComplete));
}
public CompletionStage stageToReturn(CompletionStage stage, ChannelHandlerContext ctx) {
assert ctx.channel().eventLoop().inEventLoop();
if (CompletionStages.isCompletedSuccessfully(stage)) {
RespRequestHandler rrh = CompletionStages.join(stage);
return rrh.myStage();
}
return stage.handleAsync(HANDLE_FAILURE, ctx.channel().eventLoop());
}
/**
* Handles ensuring that a stage TriConsumer is only invoked on the event loop. The TriConsumer can then do things
* such as writing to the underlying channel or update variables in a thread safe manner without additional
* synchronization or volatile variables.
*
* If the TriConsumer provided is null the provided stage must be completed only on the event loop of the
* provided ctx.
* This can be done via the Async methods on {@link CompletionStage} such as {@link CompletionStage#thenApplyAsync(Function, Executor)};
*
* If the returned stage was completed exceptionally, any exception is ignored, assumed to be handled in the
* biConsumer or further up the stage.
* The handlerWhenComplete is not invoked if the stage was completed exceptionally as well.
* If the triConsumer or handlerWhenComplete throw an exception when invoked the returned stage will
* complete with that exception.
*
* This method can only be invoked on the event loop for the provided ctx.
*
* @param The stage return value
* @param stage The stage that will complete. Note that if biConsumer is null this stage may only be completed on the event loop
* @param ctx The context used for this request, normally the event loop is the thing used
* @param biConsumer The consumer to be ran on
* @param handlerWhenComplete A function to map the result to which handler will be used for future requests. Only ever invoked on the event loop. May be null, if so the current handler is returned
* @return A stage that is only ever completed on the event loop that provides the new handler to use
*/
private CompletionStage stageToReturn(CompletionStage stage, ChannelHandlerContext ctx,
BiConsumer biConsumer, Function handlerWhenComplete) {
assert ctx.channel().eventLoop().inEventLoop();
// Only one or the other can be null
assert (biConsumer != null && handlerWhenComplete == null) || (biConsumer == null && handlerWhenComplete != null) :
"triConsumer was: " + biConsumer + " and handlerWhenComplete was: " + handlerWhenComplete;
if (CompletionStages.isCompletedSuccessfully(stage)) {
E result = CompletionStages.join(stage);
try {
if (handlerWhenComplete != null) {
return CompletableFuture.completedFuture(handlerWhenComplete.apply(result));
}
biConsumer.accept(result, allocatorToUse);
} catch (Throwable t) {
Resp3Response.error(t, allocator());
return myStage();
}
return myStage;
}
if (biConsumer != null) {
// Note that this method is only ever invoked in the event loop, so this whenCompleteAsync can never complete
// until this request completes, meaning the thenApply will always be invoked in the event loop as well
return CompletionStages.handleAndComposeAsync(stage, (e, t) -> {
if (t != null) {
Resp3Response.error(t, allocator());
} else {
try {
biConsumer.accept(e, allocatorToUse);
} catch (Throwable innerT) {
return CompletableFuture.failedFuture(innerT);
}
}
return myStage;
}, ctx.channel().eventLoop());
}
return stage.handleAsync((value, t) -> {
if (t != null) {
Resp3Response.error(t, allocator());
// If exception, never use handler
return this;
}
return handlerWhenComplete.apply(value);
}, ctx.channel().eventLoop());
}
}