import datadog.trace.api.GlobalTracer;
import datadog.trace.api.Tracer;
import datadog.trace.context.TraceScope;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;
import org.springframework.core.task.TaskDecorator;

@Configuration
public class TracingConfiguration {

    @Bean
    public TaskDecorator datadogContextPropagatingTaskDecorator() {
        return runnable -> {
            Tracer tracer = GlobalTracer.get();
            TraceScope.Continuation continuation = tracer.captureActiveSpan();

            return () -> {
                if (continuation != null) {
                    try (TraceScope scope = continuation.activate()) {
                        runnable.run();
                    }
                } else {
                    runnable.run();
                }
            };
        };
    }
}