cassandra-trunk/test/unit/org/apache/cassandra/net/proxy/InboundProxyHandler.java
/*
 * Licensed to the Apache Software Foundation (ASF) under one
 * or more contributor license agreements.  See the NOTICE file
 * distributed with this work for additional information
 * regarding copyright ownership.  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 org.apache.cassandra.net.proxy;

import java.util.ArrayDeque;
import java.util.concurrent.TimeUnit;
import java.util.function.Function;

import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelInboundHandlerAdapter;
import io.netty.util.concurrent.EventExecutor;
import io.netty.util.concurrent.ScheduledFuture;

public class InboundProxyHandler extends ChannelInboundHandlerAdapter
{
    private final ArrayDeque<Forward> forwardQueue;
    private ScheduledFuture scheduled = null;
    private final Controller controller;
    public InboundProxyHandler(Controller controller)
    {
        this.controller = controller;
        this.forwardQueue = new ArrayDeque<>(1024);
    }

    @Override
    public void channelActive(ChannelHandlerContext ctx) throws Exception
    {
        super.channelActive(ctx);
        ctx.read();
    }

    @Override
    public void channelInactive(ChannelHandlerContext ctx) throws Exception {
        controller.onDisconnect.run();

        if (scheduled != null)
        {
            scheduled.cancel(true);
            scheduled = null;
        }

        if (!forwardQueue.isEmpty())
            forwardQueue.clear();

        super.channelInactive(ctx);
    }


    @Override
    public void channelRead(ChannelHandlerContext ctx, Object msg)
    {
        Forward forward = controller.forwardStrategy.forward(ctx, msg);
        forwardQueue.offer(forward);
        maybeScheduleNext(ctx.channel().eventLoop());
        controller.onRead.run();
        ctx.channel().read();
    }

    private void maybeScheduleNext(EventExecutor executor)
    {
        if (forwardQueue.isEmpty())
        {
            // Ran out of items to process
            scheduled = null;
        }
        else if (scheduled == null)
        {
            // Schedule next available or let the last in line schedule it
            Forward forward = forwardQueue.poll();
            scheduled = forward.schedule(executor);
            scheduled.addListener((e) -> {
                scheduled = null;
                maybeScheduleNext(executor);
            });
        }
    }

    private static class Forward
    {
        final long arrivedAt;
        final long latency;
        final Runnable handler;

        private Forward(long arrivedAt, long latency, Runnable handler)
        {
            this.arrivedAt = arrivedAt;
            this.latency = latency;
            this.handler = handler;
        }

        ScheduledFuture schedule(EventExecutor executor)
        {
            long now = System.currentTimeMillis();
            long elapsed = now - arrivedAt;
            long runIn = latency - elapsed;

            if (runIn > 0)
                return executor.schedule(handler, runIn, TimeUnit.MILLISECONDS);
            else
                return executor.schedule(handler, 0, TimeUnit.MILLISECONDS);
        }
    }

    private static class ForwardNormally implements ForwardStrategy
    {
        static ForwardNormally instance = new ForwardNormally();

        public Forward forward(ChannelHandlerContext ctx, Object msg)
        {
            return new Forward(System.currentTimeMillis(),
                               0,
                               () -> ctx.fireChannelRead(msg));
        }
    }

    public interface ForwardStrategy
    {
        public Forward forward(ChannelHandlerContext ctx, Object msg);
    }

    private static class ForwardWithLatency implements ForwardStrategy
    {
        private final long latency;
        private final TimeUnit timeUnit;

        ForwardWithLatency(long latency, TimeUnit timeUnit)
        {
            this.latency = latency;
            this.timeUnit = timeUnit;
        }

        public Forward forward(ChannelHandlerContext ctx, Object msg)
        {
            return new Forward(System.currentTimeMillis(),
                               timeUnit.toMillis(latency),
                               () -> ctx.fireChannelRead(msg));
        }
    }

    private static class CloseAfterRead implements ForwardStrategy
    {
        private final Runnable afterClose;

        CloseAfterRead(Runnable afterClose)
        {
            this.afterClose = afterClose;
        }

        public Forward forward(ChannelHandlerContext ctx, Object msg)
        {
            return  new Forward(System.currentTimeMillis(),
                                0,
                                () -> {
                                    ctx.channel().close().syncUninterruptibly();
                                    afterClose.run();
                                });
        }
    }

    private static class TransformPayload<T> implements ForwardStrategy
    {
        private final Function<T, T> fn;

        TransformPayload(Function<T, T> fn)
        {
            this.fn = fn;
        }

        public Forward forward(ChannelHandlerContext ctx, Object msg)
        {
            return new Forward(System.currentTimeMillis(),
                               0,
                               () -> ctx.fireChannelRead(fn.apply((T) msg)));
        }
    }

    public static class Controller
    {
        private volatile InboundProxyHandler.ForwardStrategy forwardStrategy;
        private volatile Runnable onRead = () -> {};
        private volatile Runnable onDisconnect = () -> {};

        public Controller()
        {
            this.forwardStrategy = ForwardNormally.instance;
        }
        public void onRead(Runnable onRead)
        {
            this.onRead = onRead;
        }

        public void onDisconnect(Runnable onDisconnect)
        {
            this.onDisconnect = onDisconnect;
        }

        public void reset()
        {
            this.forwardStrategy = ForwardNormally.instance;
        }

        public void withLatency(long latency, TimeUnit timeUnit)
        {
            this.forwardStrategy = new ForwardWithLatency(latency, timeUnit);
        }

        public void withCloseAfterRead(Runnable afterClose)
        {
            this.forwardStrategy = new CloseAfterRead(afterClose);
        }

        public <T> void withPayloadTransform(Function<T, T> fn)
        {
            this.forwardStrategy = new TransformPayload<>(fn);
        }
    }

}
