com.turbospaces.grpc.ContextPreservingClientCall Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of bootstrap-grpc-core Show documentation
There is a newer version: 2.0.33
package com.turbospaces.grpc;

import java.util.Map;
import java.util.Objects;

import org.apache.commons.lang3.StringUtils;
import org.apache.commons.lang3.time.StopWatch;
import org.slf4j.MDC;

import com.turbospaces.boot.Bootstrap;
import com.turbospaces.boot.BootstrapAware;
import com.turbospaces.common.PlatformUtil;
import com.turbospaces.executor.PlatformExecutorService;
import com.turbospaces.mdc.MdcTags;

import io.grpc.ClientCall;
import io.grpc.Context;
import io.grpc.ForwardingClientCall.SimpleForwardingClientCall;
import io.grpc.ForwardingClientCallListener.SimpleForwardingClientCallListener;
import io.grpc.Metadata;
import io.grpc.MethodDescriptor;
import lombok.extern.slf4j.Slf4j;

@Slf4j
public class ContextPreservingClientCall extends SimpleForwardingClientCall implements BootstrapAware {
    private final MethodDescriptor method;
    private Bootstrap bootstrap;

    public ContextPreservingClientCall(ClientCall delegate, MethodDescriptor method) {
        super(delegate);
        this.method = Objects.requireNonNull(method);
    }
    @Override
    public void setBootstrap(Bootstrap bootstrap) {
        this.bootstrap = bootstrap;
    }
    @Override
    public void start(Listener responseListener, Metadata meta) {
        Context current = Context.current();
        Map mdc = MDC.getCopyOfContextMap();
        String messageId = Objects.requireNonNull(GrpcHeaders.CONTEXT_MESSAGE_ID.get(), "Context is not injected");
        putMetaAndDetachContext(current, meta);

        log.debug("about to make grpc call (m-{}):({})", messageId, method.getFullMethodName());
        String operation = PlatformUtil.toLowerUnderscore(method.getFullMethodName());
        StopWatch stopWatch = StopWatch.createStarted();

        super.start(new SimpleForwardingClientCallListener<>(responseListener) {
            @Override
            public void onMessage(RespT message) {
                try {
                    PlatformExecutorService.propagete(mdc);
                    try {
                        super.onMessage(message);
                    } finally {
                        log.debug("completed grpc call (m-{}):({}) in {}, body: {}", messageId, operation, stopWatch, message);
                    }
                } finally {
                    MDC.clear();
                }
            }
        }, meta);
    }
    private void putMetaAndDetachContext(Context current, Metadata meta) {
        String timestamp = Long.toString(System.currentTimeMillis());
        String messageId = GrpcHeaders.CONTEXT_MESSAGE_ID.get();
        String traceId = Objects.isNull(MDC.get(MdcTags.MDC_TRACE_ID)) ? messageId : MDC.get(MdcTags.MDC_TRACE_ID);
        String topic = bootstrap.props().CLOUD_APP_ID.get();
        String routingKey = GrpcHeaders.CONTEXT_ROUTING_KEY.get();

        meta.put(GrpcHeaders.X_MESSAGE_ID, messageId);
        meta.put(GrpcHeaders.X_TRACE_ID, traceId);
        meta.put(GrpcHeaders.X_TIMESTAMP, timestamp);
        meta.put(GrpcHeaders.X_TOPIC, topic);

        //
        // ~ optional
        //
        if (StringUtils.isNotEmpty(routingKey)) {
            meta.put(GrpcHeaders.X_ROUTING_KEY, routingKey.toString());
        }

        current.detach(Context.ROOT);
    }
}