RequestIdFilter.java

package com.sagedevices.platform.jersey.filters;

import com.google.common.base.Strings;
import com.sagedevices.platform.jersey.core.RequestId;
import com.sagedevices.platform.util.EncodingUtil;
import org.slf4j.MDC;

import javax.servlet.*;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.IOException;
import java.util.UUID;
import java.util.concurrent.ThreadLocalRandom;
import java.util.concurrent.atomic.AtomicInteger;

public class RequestIdFilter implements Filter {

    protected RequestId.GeneratorType requestIdGenerator;

    public AtomicInteger originCounter;
    public AtomicInteger localCounter;

    public RequestIdFilter() {
    }

    @Override
    public void init(FilterConfig filterConfig) throws ServletException {
        requestIdGenerator = RequestId.GeneratorType.valueOf(filterConfig.getInitParameter("RequestIdGenerator"));
        if (requestIdGenerator == RequestId.GeneratorType.SimpleIncrement) {
            originCounter = new AtomicInteger(0);
            localCounter = new AtomicInteger(0);
        }
    }

    @Override
    public void destroy() { /* unused */ }

    @Override
    public void doFilter(ServletRequest request, ServletResponse response, FilterChain chain) throws IOException, ServletException {

        HttpServletRequest req = (HttpServletRequest) request;
        HttpServletResponse resp = (HttpServletResponse) response;

        // Generate the local request id
        final String localRequestId = getNewRequestId(requestIdGenerator, localCounter);
        MDC.put(RequestId.MDC_NAME_LOCAL, localRequestId);
        resp.addHeader(RequestId.HEADER_NAME_LOCAL, EncodingUtil.urlEncode(localRequestId));

        String originRequestId = req.getHeader(RequestId.HEADER_NAME_ORIGIN);
        if (Strings.isNullOrEmpty(originRequestId)) {
            originRequestId = getNewRequestId(requestIdGenerator, originCounter);
        }
        if (!Strings.isNullOrEmpty(originRequestId)) {
            MDC.put(RequestId.MDC_NAME_ORIGIN, originRequestId);
            resp.addHeader(RequestId.HEADER_NAME_ORIGIN, EncodingUtil.urlEncode(originRequestId));
        }

        try {
            chain.doFilter(request, response);
        } finally {
            MDC.remove(RequestId.MDC_NAME_LOCAL);
            MDC.remove(RequestId.MDC_NAME_ORIGIN);
        }

    }

    public static String getNewRequestId(RequestId.GeneratorType generator, AtomicInteger counter) {
        switch (generator) {
            case SimpleIncrement:
                return String.valueOf(counter.incrementAndGet());
            case RandomLong:
                return String.valueOf(ThreadLocalRandom.current().nextLong()).substring(1);
            case UUID:
            default:
                return UUID.randomUUID().toString();
        }
    }

}