octoteam.tahiti.server.pipeline.RequestRateLimitHandler.java Source code

Introduction

Here is the source code for octoteam.tahiti.server.pipeline.RequestRateLimitHandler.java

Source

package octoteam.tahiti.server.pipeline;

import io.netty.channel.ChannelHandler;
import io.netty.channel.ChannelHandlerContext;
import octoteam.tahiti.protocol.SocketMessageProtos.Message;
import octoteam.tahiti.server.event.RateLimitExceededEvent;
import octoteam.tahiti.server.session.PipelineHelper;
import octoteam.tahiti.shared.netty.ExtendedContext;
import octoteam.tahiti.shared.netty.MessageHandler;
import octoteam.tahiti.shared.protocol.ProtocolUtil;
import wheellllll.license.License;

import java.util.concurrent.Callable;

/**
 * ?????. ??,  RateLimitExceededEvent ,
 * ? LIMIT_EXCEEDED ?.
 * ??.
 */
@ChannelHandler.Sharable
public class RequestRateLimitHandler extends MessageHandler {

    private final Message.ServiceCode serviceCode;
    private final String name;
    private final String sessionKey;
    private final Callable<License> quotaLimiterFactory;

    /**
     * @param serviceCode ???, ????
     * @param name        ???
     * @param factory     ?? Callable , ????
     */
    public RequestRateLimitHandler(ExtendedContext extendedContext, Message.ServiceCode serviceCode, String name,
            Callable<License> factory) {
        super(extendedContext);
        this.serviceCode = serviceCode;
        this.name = name;
        this.sessionKey = "limiter_" + name;
        this.quotaLimiterFactory = factory;
    }

    @Override
    protected void messageReceived(ChannelHandlerContext ctx, Message msg) throws Exception {
        if (msg.getDirection() != Message.DirectionCode.REQUEST || msg.getService() != serviceCode) {
            ctx.fireChannelRead(msg);
            return;
        }
        License quotaLimiter = (License) PipelineHelper.getSession(ctx).get(sessionKey);
        if (quotaLimiter == null) {
            quotaLimiter = this.quotaLimiterFactory.call();
            PipelineHelper.getSession(ctx).put(sessionKey, quotaLimiter);
        }
        if (quotaLimiter.use() == License.Availability.AVAILABLE) {
            ctx.fireChannelRead(msg);
        } else {
            RateLimitExceededEvent evt = new RateLimitExceededEvent(name);
            ctx.fireUserEventTriggered(evt);
            ctx.writeAndFlush(buildExceededMsg(msg));
        }
    }

    /**
     * @param req ???
     * @return message ? LIMIT_EXCEEDED ?
     */
    private Message buildExceededMsg(Message req) {
        Message.Builder resp = ProtocolUtil.buildResponse(req).setStatus(Message.StatusCode.LIMIT_EXCEEDED);
        return resp.build();
    }

}