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

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

import java.util.function.Supplier;

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

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

/** {@linkplain Supplier} 包装类（注入 {@linkplain org.slf4j.MDC MDC} 和 traceId 调用链跟踪信息） */
@TraceCrossThread
public class TracingSupplierWrapper extends SupplierWrapper
{
    private final MdcAttr mdcAttr;
    
    public TracingSupplierWrapper(Supplier s)
    {
        this(s, MdcAttr.fromMdc());
    }

    public TracingSupplierWrapper(Supplier s, MdcAttr mdcAttr)
    {
        super(s);
        
        this.mdcAttr = mdcAttr;
    }

    @Override
    public V get()
    {
        try
        {
            mdcAttr.putMdc();
            
            return super.get();
        }
        finally
        {
            mdcAttr.removeMdc();
        }
    }

    public static  TracingSupplierWrapper of(Supplier s)
    {
        return new TracingSupplierWrapper<>(s);
    }
}