com.bloom.zerofs.rest.HealthCheckHandler.java Source code

Introduction

Here is the source code for com.bloom.zerofs.rest.HealthCheckHandler.java

Source

/**
 * Copyright 2016 Bloom Corp. All rights reserved.
 *
 * 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.
 */
package com.bloom.zerofs.rest;

import io.netty.buffer.Unpooled;
import io.netty.channel.ChannelDuplexHandler;
import io.netty.channel.ChannelFuture;
import io.netty.channel.ChannelFutureListener;
import io.netty.channel.ChannelHandlerContext;
import io.netty.channel.ChannelPromise;
import io.netty.handler.codec.http.DefaultFullHttpResponse;
import io.netty.handler.codec.http.FullHttpResponse;
import io.netty.handler.codec.http.HttpHeaders;
import io.netty.handler.codec.http.HttpRequest;
import io.netty.handler.codec.http.HttpResponseStatus;
import io.netty.handler.codec.http.HttpVersion;
import io.netty.handler.codec.http.LastHttpContent;
import io.netty.util.ReferenceCountUtil;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

/**
 * Responsible for responding to health check requests
 * {@link RestServerState} assists in knowing the state of the system at any point in time
 */
public class HealthCheckHandler extends ChannelDuplexHandler {
    private static final byte[] GOOD = "GOOD".getBytes();
    private static final byte[] BAD = "BAD".getBytes();

    private final String healthCheckUri;
    private final RestServerState restServerState;
    private final NettyMetrics nettyMetrics;
    private final Logger logger = LoggerFactory.getLogger(getClass());

    private HttpRequest request;
    private FullHttpResponse response;
    private long startTimeInMs;

    public HealthCheckHandler(RestServerState restServerState, NettyMetrics nettyMetrics) {
        this.restServerState = restServerState;
        this.healthCheckUri = restServerState.getHealthCheckUri();
        this.nettyMetrics = nettyMetrics;
        logger.trace("Created HealthCheckHandler for HealthCheckUri=" + healthCheckUri);
    }

    @Override
    public void channelRead(ChannelHandlerContext ctx, Object obj) throws Exception {
        logger.trace("Reading on channel {}", ctx.channel());
        boolean forwardObj = false;
        if (obj instanceof HttpRequest) {
            if (request == null && ((HttpRequest) obj).getUri().equals(healthCheckUri)) {
                nettyMetrics.healthCheckRequestRate.mark();
                startTimeInMs = System.currentTimeMillis();
                logger.trace("Handling health check request while in state " + restServerState.isServiceUp());
                request = (HttpRequest) obj;
                if (restServerState.isServiceUp()) {
                    response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1, HttpResponseStatus.OK,
                            Unpooled.wrappedBuffer(GOOD));
                    HttpHeaders.setKeepAlive(response, HttpHeaders.isKeepAlive(request));
                    HttpHeaders.setContentLength(response, GOOD.length);
                } else {
                    response = new DefaultFullHttpResponse(HttpVersion.HTTP_1_1,
                            HttpResponseStatus.SERVICE_UNAVAILABLE, Unpooled.wrappedBuffer(BAD));
                    HttpHeaders.setKeepAlive(response, false);
                    HttpHeaders.setContentLength(response, BAD.length);
                }
                nettyMetrics.healthCheckRequestProcessingTimeInMs
                        .update(System.currentTimeMillis() - startTimeInMs);
            } else {
                // Rest server could be down even if not for health check request. We intentionally don't take any action in this
                // handler for such cases and leave it to the downstream handlers to handle it
                forwardObj = true;
            }
        }
        if (obj instanceof LastHttpContent) {
            if (response != null) {
                // response was created when we received the request with health check uri
                ChannelFuture future = ctx.writeAndFlush(response);
                if (!HttpHeaders.isKeepAlive(response)) {
                    future.addListener(ChannelFutureListener.CLOSE);
                }
                request = null;
                response = null;
                nettyMetrics.healthCheckRequestRoundTripTimeInMs.update(System.currentTimeMillis() - startTimeInMs);
            } else {
                // request was not for health check uri
                forwardObj = true;
            }
        } else if (request == null) {
            // http Content which is not LastHttpContent is not intended for this handler
            forwardObj = true;
        }
        if (forwardObj) {
            super.channelRead(ctx, obj);
        } else {
            ReferenceCountUtil.release(obj);
        }
    }

    @Override
    public void write(ChannelHandlerContext ctx, Object msg, ChannelPromise promise) throws Exception {
        if (!restServerState.isServiceUp()) {
            if (msg instanceof LastHttpContent) {
                // Start closing client channels after we've completed writing to them (even if they are keep-alive)
                logger.info("Health check request handler closing connection " + ctx.channel()
                        + " since in shutdown mode.");
                promise.addListener(ChannelFutureListener.CLOSE);
                nettyMetrics.healthCheckHandlerChannelCloseOnWriteCount.inc();
            }
        }
        super.write(ctx, msg, promise);
    }
}