TransmittableThreadLocal正确使用方式

时间:2025-03-30 15:16:57
/** * Capture all {@link TransmittableThreadLocal} values in current thread. * * @return the captured {@link TransmittableThreadLocal} values * @since 2.3.0 */ @Nonnull public static Object capture() { Map<TransmittableThreadLocal<?>, Object> captured = new HashMap<TransmittableThreadLocal<?>, Object>(); for (TransmittableThreadLocal<?> threadLocal : holder.get().keySet()) { captured.put(threadLocal, threadLocal.copyValue()); } return captured; } /** * Replay the captured {@link TransmittableThreadLocal} values from {@link #capture()}, * and return the backup {@link TransmittableThreadLocal} values in current thread before replay. * * @param captured captured {@link TransmittableThreadLocal} values from other thread from {@link #capture()} * @return the backup {@link TransmittableThreadLocal} values before replay * @see #capture() * @since 2.3.0 */ @Nonnull public static Object replay(@Nonnull Object captured) { @SuppressWarnings("unchecked") Map<TransmittableThreadLocal<?>, Object> capturedMap = (Map<TransmittableThreadLocal<?>, Object>) captured; Map<TransmittableThreadLocal<?>, Object> backup = new HashMap<TransmittableThreadLocal<?>, Object>(); for (Iterator<? extends Map.Entry<TransmittableThreadLocal<?>, ?>> iterator = holder.get().entrySet().iterator(); iterator.hasNext(); ) { Map.Entry<TransmittableThreadLocal<?>, ?> next = iterator.next(); TransmittableThreadLocal<?> threadLocal = next.getKey(); // backup backup.put(threadLocal, threadLocal.get()); // clear the TTL values that is not in captured // avoid the extra TTL values after replay when run task if (!capturedMap.containsKey(threadLocal)) { iterator.remove(); threadLocal.superRemove(); } } // set values to captured TTL setTtlValuesTo(capturedMap); // call beforeExecute callback doExecuteCallback(true); return backup; } /** * Clear all {@link TransmittableThreadLocal} values in current thread, * and return the backup {@link TransmittableThreadLocal} values in current thread before clear. * * @return the backup {@link TransmittableThreadLocal} values before clear * @since 2.9.0 */ @Nonnull public static Object clear() { return replay(Collections.emptyMap()); } /** * Restore the backup {@link TransmittableThreadLocal} values from {@link #replay(Object)}/{@link #clear()}. * * @param backup the backup {@link TransmittableThreadLocal} values from {@link #replay(Object)}/{@link #clear()} * @see #replay(Object) * @see #clear() * @since 2.3.0 */ public static void restore(@Nonnull Object backup) { @SuppressWarnings("unchecked") Map<TransmittableThreadLocal<?>, Object> backupMap = (Map<TransmittableThreadLocal<?>, Object>) backup; // call afterExecute callback doExecuteCallback(false); for (Iterator<? extends Map.Entry<TransmittableThreadLocal<?>, ?>> iterator = holder.get().entrySet().iterator(); iterator.hasNext(); ) { Map.Entry<TransmittableThreadLocal<?>, ?> next = iterator.next(); TransmittableThreadLocal<?> threadLocal = next.getKey(); // clear the TTL values that is not in backup // avoid the extra TTL values after restore if (!backupMap.containsKey(threadLocal)) { iterator.remove(); threadLocal.superRemove(); } } // restore TTL values setTtlValuesTo(backupMap); } private static void setTtlValuesTo(@Nonnull Map<TransmittableThreadLocal<?>, Object> ttlValues) { for (Map.Entry<TransmittableThreadLocal<?>, Object> entry : ttlValues.entrySet()) { @SuppressWarnings("unchecked") TransmittableThreadLocal<Object> threadLocal = (TransmittableThreadLocal<Object>) entry.getKey(); threadLocal.set(entry.getValue()); } }