com.proofpoint.tracetoken.TraceTokenManager Maven / Gradle / Ivy
/*
* Copyright 2010 Proofpoint, Inc.
*
* 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
*
* http://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 com.proofpoint.tracetoken;
import com.google.auto.value.AutoValue;
import com.google.common.collect.ImmutableMap;
import jakarta.annotation.Nullable;
import java.util.LinkedHashMap;
import java.util.Map;
import java.util.UUID;
import static com.google.common.base.Preconditions.checkArgument;
import static java.lang.Thread.currentThread;
import static java.util.Objects.requireNonNull;
/**
* Utility class for managing the trace token, a request identifier and set of
* properties associated with a thread while the thread is handling the request.
*/
public final class TraceTokenManager
{
private static final ThreadLocal token = new ThreadLocal<>();
private TraceTokenManager()
{}
/**
* Associate a given trace token id, with no other properties, with the
* current thread.
*
* @param tokenId The tokenId to associate with the current thread, or null to
* remove the thread's token.
* @return a {@link TraceTokenScope} which may be used to restore the thread's
* previous association. Intended to be used with try-with-resources:
*
* try (TraceTokenScope ignored = registerRequestToken(traceTokenId)) {
* // process request
* }
*
*/
public static TraceTokenScope registerRequestToken(@Nullable String tokenId)
{
if (tokenId == null) {
return registerTraceToken(null);
}
return registerTraceToken(new TraceToken(ImmutableMap.of("id", tokenId)));
}
/**
* Associate a given trace token with the current thread.
*
* @param token The {@link TraceToken} to associate with the current thread, or null to
* remove the thread's token.
* @return a {@link TraceTokenScope} which may be used to restore the thread's
* previous association. Intended to be used with try-with-resources:
*
* try (TraceTokenScope ignored = registerTraceToken(traceToken)) {
* // process request
* }
*
*/
public static TraceTokenScope registerTraceToken(@Nullable TraceToken token)
{
TokenState oldTokenState = TraceTokenManager.token.get();
String oldThreadName;
if (oldTokenState == null) {
oldThreadName = currentThread().getName();
}
else {
oldThreadName = oldTokenState.getOldThreadName();
}
if (token == null) {
TraceTokenManager.token.set(null);
currentThread().setName(oldThreadName);
}
else {
TraceTokenManager.token.set(new AutoValue_TraceTokenManager_TokenState(token, oldThreadName));
currentThread().setName(oldThreadName + " " + token);
}
if (oldTokenState == null) {
return new TraceTokenScope(null);
}
else {
return new TraceTokenScope(oldTokenState.getToken());
}
}
/**
* @return The current thread's trace token in string form, or null if no token.
* @deprecated Use {@link #getCurrentTraceToken()}.
*/
@Deprecated
@Nullable
public static String getCurrentRequestToken()
{
TokenState tokenState = token.get();
if (tokenState == null) {
return null;
}
return tokenState.getToken().toString();
}
/**
* @return The current thread's trace token, or null if no token.
*/
@Nullable
public static TraceToken getCurrentTraceToken()
{
TokenState tokenState = token.get();
if (tokenState == null) {
return null;
}
return tokenState.getToken();
}
/**
* Create and register a new trace token with a random id.
* @param properties Additional properties to include in the token.
* @return The id of the created token.
*/
public static String createAndRegisterNewRequestToken(String... properties)
{
checkArgument((properties.length % 2) == 0, "odd number of elements in properties");
String newToken = UUID.randomUUID().toString();
registerRequestToken(newToken);
if (properties.length != 0) {
addTraceTokenProperties(properties);
}
return newToken;
}
/**
* Remove the thread's token.
*/
public static void clearRequestToken()
{
TokenState oldToken = TraceTokenManager.token.get();
token.remove();
if (oldToken != null) {
currentThread().setName(oldToken.getOldThreadName());
}
}
/**
* Add properties to the current thread's trace token. If there is
* currently no trace token, does nothing.
*
* @param properties Properties to add or replace.
* @return a {@link TraceTokenScope} which may be used to restore the thread's
* previous set of properties.
*/
public static TraceTokenScope addTraceTokenProperties(String... properties)
{
TokenState tokenState = token.get();
if (tokenState == null) {
return new TraceTokenScope(null);
}
Map map = new LinkedHashMap<>(tokenState.getToken());
checkArgument((properties.length % 2) == 0, "odd number of elements in properties");
for (int i = 0; i < properties.length; i += 2) {
requireNonNull(properties[i], "property key is null");
requireNonNull(properties[i+1], "property value is null");
map.put(properties[i], properties[i+1]);
}
return registerTraceToken(new TraceToken(map));
}
@AutoValue
abstract static class TokenState
{
abstract TraceToken getToken();
abstract String getOldThreadName();
}
}