io.github.hpsocket.soa.starter.skywalking.async.TracingFunctionWrapper Maven / Gradle / Ivy

Go to download
package io.github.hpsocket.soa.starter.skywalking.async;

import java.util.function.Function;

import org.apache.skywalking.apm.toolkit.trace.FunctionWrapper;
import org.apache.skywalking.apm.toolkit.trace.TraceCrossThread;

import io.github.hpsocket.soa.framework.core.mdc.MdcAttr;

/** {@linkplain Function} 包装类（注入 {@linkplain org.slf4j.MDC MDC} 和 traceId 调用链跟踪信息） */
@TraceCrossThread
public class TracingFunctionWrapper extends FunctionWrapper
{
    private final MdcAttr mdcAttr;
    
    public TracingFunctionWrapper(Function f)
    {
        this(f, MdcAttr.fromMdc());
    }
    
    public TracingFunctionWrapper(Function f, MdcAttr mdcAttr)
    {
        super(f);
        
        this.mdcAttr = mdcAttr;
    }
    
    @Override
    public R apply(T t)
    {
        try
        {
            mdcAttr.putMdc();
            
            return super.apply(t);
        }
        finally
        {
            mdcAttr.removeMdc();
        }
    }

    public static  TracingFunctionWrapper of(Function f)
    {
        return new TracingFunctionWrapper<>(f);
    }

}