io.reactivex.netty.contexts.ThreadLocalRequestCorrelator Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of rxnetty-contexts Show documentation
rxnetty-contexts
There is a newer version: 0.5.1
/*
 * Copyright 2014 Netflix, Inc.
 *
 * 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.reactivex.netty.contexts;

import io.netty.channel.EventLoopGroup;
import io.netty.util.AttributeKey;
import io.netty.util.AttributeMap;
import io.netty.util.DefaultAttributeMap;
import io.reactivex.netty.client.RxClient;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import java.io.PrintStream;
import java.util.Stack;
import java.util.concurrent.Callable;

/**
 * An implementation of {@link RequestCorrelator} that fetches {@link ContextsContainer} and request Identifer from
 * {@link ThreadLocal} variables. 

 * The usage of this correlator assumes that every {@link EventLoopGroup} used by all {@link RxClient}s in use for
 * the propagation of these contexts, uses {@link ContextAwareEventLoopGroup}. 

 * In case, the application is using any other threadpools, the invariant of the thread state copied from one thread
 * to another is maintained.
 *
 * @author Nitesh Kant
 */
public class ThreadLocalRequestCorrelator implements RequestCorrelator {

    private static final Logger logger = LoggerFactory.getLogger(ThreadLocalRequestCorrelator.class);

    protected static final ThreadLocal state = new ThreadLocal() {
        @Override
        protected ThreadStateHolder initialValue() {
            return new ThreadStateHolder();
        }
    };

    @Override
    public String getRequestIdForClientRequest() {
        return state.get().getRequestId();
    }

    @Override
    public ContextsContainer getContextForClientRequest(String requestId) {
        return state.get().getContainer();
    }

    @Override
    public void onNewServerRequest(String requestId, ContextsContainer contextsContainer) {
        if (null == requestId) {
            throw new IllegalArgumentException("Request Id can not be null.");
        }
        if (null == contextsContainer) {
            throw new IllegalArgumentException("Context container can not be null.");
        }

        state.get().push(new DefaultAttributeMap());

        state.get().setRequestId(requestId);
        state.get().setContainer(contextsContainer);
    }

    @Override
    public void beforeNewClientRequest(String requestId, ContextsContainer contextsContainer) {
        onNewServerRequest(requestId, contextsContainer); // same effect
    }

    @Override
    public void onClientProcessingEnd(String requestId) {
        state.get().pop();
    }

    @Override
    public void onServerProcessingEnd(String requestId) {
        state.get().pop();
    }

    public static ContextsContainer getCurrentContextContainer() {
        return state.get().getContainer();
    }

    public static String getCurrentRequestId() {
        return state.get().getRequestId();
    }

    @Override
    public  Callable makeClosure(final Callable original) {

        final AttributeMap currentState = state.get().peek();

        return new Callable() {
            @Override
            public V call() throws Exception {
                state.get().push(currentState);
                try {
                    return original.call();
                } finally {
                    state.get().pop();
                }
            }
        };
    }

    @Override
    public Runnable makeClosure(final Runnable original) {
        final AttributeMap currentState = state.get().peek();

        return new Runnable() {

            @Override
            public void run() {
                state.get().push(currentState);
                try {
                    original.run();
                } finally {
                    state.get().pop();
                }
            }
        };
    }

    public void dumpThreadState() {
        if (logger.isDebugEnabled()) {
            ThreadStateHolder threadStateHolder = state.get();
            logger.debug("******************** Thread Attributes ********************");
            int count = 0;
            for (AttributeMap threadAttribute : threadStateHolder.threadAttributes) {
                logger.debug("Stack level==> " + ++count);
                logger.debug("Request Id: " + threadAttribute.attr(ThreadStateHolder.requestIdKey));
                logger.debug("Context container: " + threadAttribute.attr(ThreadStateHolder.containerKey));
            }
            logger.debug("***********************************************************");
        }
    }

    public void dumpThreadState(PrintStream outputStream) {
        ThreadStateHolder threadStateHolder = state.get();
        outputStream.println("******************** Thread Attributes ********************");
        int count = 0;
        for (AttributeMap threadAttribute : threadStateHolder.threadAttributes) {
            outputStream.println("Stack level==> " + ++count);
            outputStream.println("Request Id: " + threadAttribute.attr(ThreadStateHolder.requestIdKey));
            outputStream.println("Context container: " + threadAttribute.attr(ThreadStateHolder.containerKey));
        }
        outputStream.println("***********************************************************");
    }

    public static final class ThreadStateHolder {

        private static final AttributeKey requestIdKey =
                AttributeKey.valueOf("rxnetty-contexts-threadlocal-request-id-key");
        private static final AttributeKey containerKey =
                AttributeKey.valueOf("rxnetty-contexts-threadlocal-context-container-key");

        /**
         * This is a stack because the client and server processing can be on the same thread, either by chance or
         * just because there is only a single thread in the system.
         * This means we can not simply set & unset on start & stop of processing as it can step on other's feet.
         * The best way to handle is to use a stack where in states are pushed and popped. If we pop everything, it is
         * as good as clear state.
         * Here {@link AttributeMap} is used just to have a map that stores multiple objects and still is type safe.
         */
        private final Stack threadAttributes = new Stack();

        public void push(AttributeMap map) {
            threadAttributes.push(map);
        }

        public AttributeMap pop() {
            return threadAttributes.pop();
        }

        public AttributeMap peek() {
            if (isEmpty()) {
                return null;
            }
            return threadAttributes.peek();
        }

        public boolean isEmpty() {
            return threadAttributes.isEmpty();
        }

        public void setRequestId(String requestId) {
            if (!isEmpty()) {
                peek().attr(requestIdKey).set(requestId);
            }
        }

        public void setContainer(ContextsContainer container) {
            if (!isEmpty()) {
                peek().attr(containerKey).set(container);
            }
        }

        public String getRequestId() {
            if (isEmpty()) {
                return null;
            }
            return peek().attr(requestIdKey).get();
        }

        public ContextsContainer getContainer() {
            if (isEmpty()) {
                return null;
            }
            return peek().attr(containerKey).get();
        }
    }
}