io.atomix.catalyst.transport.netty.NettyHandler.java Source code

Introduction

Here is the source code for io.atomix.catalyst.transport.netty.NettyHandler.java

Source

/*
 * Copyright 2015 the original author or authors.
 *
 * Licensed 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 io.atomix.catalyst.transport.netty;

import io.atomix.catalyst.concurrent.SingleThreadContext;
import io.atomix.catalyst.concurrent.ThreadContext;
import io.atomix.catalyst.transport.Connection;
import io.netty.buffer.ByteBuf;
import io.netty.channel.Channel;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.handler.timeout.IdleState;
import io.netty.handler.timeout.IdleStateEvent;

import java.util.Map;
import java.util.concurrent.CompletableFuture;
import java.util.function.Consumer;

/**
 * Netty handler.
 *
 * @author <a href="http://github.com/kuujo">Jordan Halterman</a>
 */
class NettyHandler extends ChannelInboundHandlerAdapter {
    private final Map<Channel, NettyConnection> connections;
    private final Consumer<Connection> listener;
    private final ThreadContext context;
    private final NettyOptions options;

    protected NettyHandler(Map<Channel, NettyConnection> connections, Consumer<Connection> listener,
            ThreadContext context, NettyOptions options) {
        this.connections = connections;
        this.listener = listener;
        this.context = context;
        this.options = options;
    }

    /**
     * Adds a connection for the given channel.
     *
     * @param channel The channel for which to add the connection.
     * @param connection The connection to add.
     */
    protected void setConnection(Channel channel, NettyConnection connection) {
        connections.put(channel, connection);
    }

    /**
     * Returns the connection for the given channel.
     *
     * @param channel The channel for which to return the connection.
     * @return The connection.
     */
    protected NettyConnection getConnection(Channel channel) {
        return connections.get(channel);
    }

    /**
     * Removes the connection for the given channel.
     *
     * @param channel The channel for which to remove the connection.
     * @return The connection.
     */
    protected NettyConnection removeConnection(Channel channel) {
        return connections.remove(channel);
    }

    /**
     * Returns the current execution context or creates one.
     */
    private ThreadContext getOrCreateContext(Channel channel) {
        ThreadContext context = ThreadContext.currentContext();
        if (context != null) {
            return context;
        }
        return new SingleThreadContext(Thread.currentThread(), channel.eventLoop(),
                this.context.serializer().clone());
    }

    @Override
    public void channelActive(ChannelHandlerContext context) throws Exception {
        Channel channel = context.channel();
        NettyConnection connection = new NettyConnection(channel, getOrCreateContext(channel), options);
        setConnection(channel, connection);
        // synchronously notify listeners in order to ensure message handlers
        // are registered before messages are handled.
        CompletableFuture.runAsync(() -> listener.accept(connection), this.context.executor()).join();
    }

    @Override
    public void channelRead(final ChannelHandlerContext context, Object message) {
        ByteBuf buffer = (ByteBuf) message;
        int type = buffer.readByte();
        switch (type) {
        case NettyConnection.REQUEST:
            handleRequest(buffer, context);
            break;
        case NettyConnection.RESPONSE:
            handleResponse(buffer, context);
            break;
        }
    }

    /**
     * Handles a request.
     */
    private void handleRequest(ByteBuf request, ChannelHandlerContext context) {
        NettyConnection connection = getConnection(context.channel());
        if (connection != null) {
            connection.handleRequest(request);
        }
    }

    /**
     * Handles a response.
     */
    private void handleResponse(ByteBuf response, ChannelHandlerContext context) {
        NettyConnection connection = getConnection(context.channel());
        if (connection != null) {
            connection.handleResponse(response);
        }
    }

    @Override
    public void exceptionCaught(ChannelHandlerContext context, final Throwable t) throws Exception {
        Channel channel = context.channel();
        NettyConnection connection = getConnection(channel);
        if (connection != null) {
            try {
                if (channel.isOpen()) {
                    channel.close();
                }
            } catch (Throwable ignore) {
            }
            connection.handleException(t);
        } else {
            channel.close();
        }
    }

    @Override
    public void channelInactive(ChannelHandlerContext context) throws Exception {
        Channel channel = context.channel();
        NettyConnection connection = removeConnection(channel);
        if (connection != null) {
            connection.handleClosed();
        }
    }

    @Override
    public void userEventTriggered(ChannelHandlerContext context, Object event) throws Exception {
        if (event instanceof IdleStateEvent && ((IdleStateEvent) event).state() == IdleState.ALL_IDLE) {
            context.close();
        }
    }

}