io.pyroscope.labels.ScopedContext Maven / Gradle / Ivy

Go to download
package io.pyroscope.labels;

import one.profiler.AsyncProfiler;

import java.util.Collections;
import java.util.HashMap;
import java.util.Map;

class ScopedContext implements AutoCloseable {
    static final ThreadLocal context = ThreadLocal.withInitial(() ->
            new Context(0L, Collections.emptyMap())
    );

    final Context previous;
    final Context current;
    final Ref, Ref>> currentRef;

    ScopedContext(LabelsSet labels) {
        previous = context.get();
        Map, Ref> nextContext = new HashMap<>(
                previous.labels.size() + labels.args.length / 2
        );
        for (Map.Entry, Ref> it : previous.labels.entrySet()) {
            Ref key = it.getKey();
            Ref value = it.getValue();

            assertAlive(key.refCount.get());
            assertAlive(value.refCount.get());

            assertAlive(key.refCount.incrementAndGet());
            assertAlive(value.refCount.incrementAndGet());

            nextContext.put(key, value);
        }

        for (int i = 0; i < labels.args.length; i += 2) {
            String ks = labels.args[i].toString();
            String vs = labels.args[i + 1].toString();
            Ref k = RefCounted.strings.acquireRef(ks);
            Ref v = RefCounted.strings.acquireRef(vs);

            Ref prev = nextContext.put(k, v);
            if (prev != null) {
                assertAlive(k.refCount.decrementAndGet());
                assertAlive(prev.refCount.decrementAndGet());
            }
        }

        boolean[] fresh = new boolean[1];
        currentRef = RefCounted.contexts.acquireRef(nextContext, fresh);
        if (!fresh[0]) {
            for (Map.Entry, Ref> it : nextContext.entrySet()) {
                it.getKey().refCount.decrementAndGet();
                it.getValue().refCount.decrementAndGet();
            }
        }

        AsyncProfiler.getInstance().setContextId(currentRef.id);
        current = new Context(currentRef.id, nextContext);
        context.set(current);
    }


    @Override
    public void close() {
        currentRef.refCount.decrementAndGet();
        context.set(previous);
        AsyncProfiler.getInstance().setContextId(previous.id);
    }

    static class Context {
        public final Long id;
        public final Map, Ref> labels;

        public Context(Long id, Map, Ref> labels) {
            this.id = id;
            this.labels = labels;
        }
    }

    static void assertAlive(long counter) {
        if (counter <= 0) {
            throw new AssertionError();
        }
    }
}