io.github.linyimin0812.profiler.extension.enhance.springbean.PersistentThreadLocal Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of spring-profiler-extension Show documentation

There is a newer version: 3.1.4

Show newest version

package io.github.linyimin0812.profiler.extension.enhance.springbean;

import java.util.Collection;
import java.util.Collections;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Supplier;

/**
 * @author linyimin
 **/
public class PersistentThreadLocal extends ThreadLocal {

    final Map allValues;
    final Supplier valueGetter;

    public PersistentThreadLocal(Supplier initialValue) {
        this(0, initialValue);
    }

    public PersistentThreadLocal(int numThread, Supplier initialValue) {
        allValues = Collections.synchronizedMap(
                numThread > 0 ? new HashMap<>(numThread) : new HashMap<>()
        );

        valueGetter = initialValue;
    }

    @Override
    protected T initialValue() {
        T value = valueGetter != null ? valueGetter.get() : super.initialValue();
        allValues.put(Thread.currentThread(), value);
        return value;
    }

    @Override
    public void set(T value) {
        super.set(value);
        allValues.put(Thread.currentThread(), value);
    }

    @Override
    public void remove() {
        super.remove();
        allValues.remove(Thread.currentThread());
    }

    public Collection getAll() {
        return allValues.values();
    }

    public void clear() {
        allValues.clear();
    }
}