io.micrometer.tracing.test.simple.SimpleBaggageManager Maven / Gradle / Ivy
/**
* Copyright 2023 the original author or authors.
*
* 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
*
* https://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.micrometer.tracing.test.simple;
import io.micrometer.common.lang.Nullable;
import io.micrometer.tracing.Baggage;
import io.micrometer.tracing.BaggageInScope;
import io.micrometer.tracing.BaggageManager;
import io.micrometer.tracing.TraceContext;
import java.util.Collections;
import java.util.HashSet;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
import java.util.stream.Collectors;
/**
* A test implementation of a baggage manager.
*
* @author Marcin Grzejszczak
* @author Jonatan Ivanov
* @since 1.0.0
*/
public class SimpleBaggageManager implements BaggageManager {
private final Map>> baggagesByContext = new ConcurrentHashMap<>();
private final SimpleTracer simpleTracer;
/**
* Creates a new instance of {@link SimpleBaggageManager}.
* @param simpleTracer simple tracer
*/
public SimpleBaggageManager(SimpleTracer simpleTracer) {
this.simpleTracer = simpleTracer;
}
@Override
public Map getAllBaggage() {
TraceContext context = simpleTracer.currentTraceContext().context();
if (context == null) {
return Collections.emptyMap();
}
return getAllBaggageForCtx(context);
}
Map getAllBaggageForCtx(TraceContext context) {
Map map = this.baggagesByContext
.getOrDefault(context, ThreadLocal.withInitial(Collections::emptySet))
.get()
.stream()
.filter(s -> s.get(context) != null)
.collect(Collectors.toMap(Baggage::name, baggage -> baggage.get(context)));
map.putAll(((SimpleTraceContext) context).baggageFromParent());
return map;
}
@Override
public Baggage getBaggage(String name) {
Baggage baggage = getBaggage(simpleTracer.currentTraceContext().context(), name);
return baggageOrNoop(baggage);
}
@Override
public Baggage getBaggage(TraceContext traceContext, String name) {
SimpleBaggageInScope simpleBaggageInScope = baggageForName(traceContext, name);
return baggageOrNoop(simpleBaggageInScope);
}
private static Baggage baggageOrNoop(Baggage baggage) {
if (baggage == null) {
return Baggage.NOOP;
}
return baggage;
}
@Nullable
private SimpleBaggageInScope baggageForName(TraceContext traceContext, String name) {
if (traceContext == null) {
return null;
}
return this.baggagesByContext.getOrDefault(traceContext, ThreadLocal.withInitial(Collections::emptySet))
.get()
.stream()
.filter(bag -> name.equalsIgnoreCase(bag.name()))
.findFirst()
.orElse(null);
}
@Override
@Deprecated
public Baggage createBaggage(String name) {
return createSimpleBaggage(name);
}
private Baggage createSimpleBaggage(String name) {
TraceContext current = simpleTracer.currentTraceContext().context();
SimpleBaggageInScope baggage = baggageForName(current, name);
if (baggage == null) {
ThreadLocal> baggages = this.baggagesByContext.getOrDefault(current,
ThreadLocal.withInitial(HashSet::new));
baggage = new SimpleBaggageInScope(simpleTracer.currentTraceContext(), name, current);
baggages.get().add(baggage);
this.baggagesByContext.put(current, baggages);
}
return baggage;
}
@Override
@Deprecated
public Baggage createBaggage(String name, String value) {
Baggage baggage = createSimpleBaggage(name);
baggage.set(value);
return baggage;
}
@Override
public BaggageInScope createBaggageInScope(String name, String value) {
return createSimpleBaggage(name).makeCurrent(value);
}
@Override
public BaggageInScope createBaggageInScope(TraceContext traceContext, String name, String value) {
return createSimpleBaggage(name).makeCurrent(traceContext, value);
}
}