org.opensearch.migrations.replay.util.OrderedWorkerTracker Maven / Gradle / Ivy

Go to download

Show more of this group Show more artifacts with this name
Show all versions of trafficReplayer Show documentation

Everything opensearch migrations

There is a newer version: 0.2.0.4

package org.opensearch.migrations.replay.util;

import java.util.Comparator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentSkipListSet;
import java.util.stream.Stream;

import org.opensearch.migrations.replay.TrafficReplayerTopLevel;
import org.opensearch.migrations.replay.datatypes.UniqueReplayerRequestKey;

import lombok.AllArgsConstructor;
import lombok.Getter;

public class OrderedWorkerTracker implements TrafficReplayerTopLevel.IStreamableWorkTracker {
    @AllArgsConstructor
    static class TimeKeyAndFuture {
        @Getter
        final long nanoTimeKey;
        final TrackedFuture future;
    }

    ConcurrentHashMap> primaryMap = new ConcurrentHashMap<>();
    ConcurrentSkipListSet> orderedSet = new ConcurrentSkipListSet<>(
        Comparator.comparingLong(TimeKeyAndFuture::getNanoTimeKey).thenComparingLong(System::identityHashCode)
    );

    @Override
    public void put(UniqueReplayerRequestKey uniqueReplayerRequestKey, TrackedFuture completableFuture) {
        var timedValue = new TimeKeyAndFuture<>(System.nanoTime(), completableFuture);
        primaryMap.put(uniqueReplayerRequestKey, timedValue);
        orderedSet.add(timedValue);
    }

    @Override
    public void remove(UniqueReplayerRequestKey uniqueReplayerRequestKey) {
        var timedValue = primaryMap.remove(uniqueReplayerRequestKey);
        assert timedValue != null;
        orderedSet.remove(timedValue);
    }

    @Override
    public boolean isEmpty() {
        return primaryMap.isEmpty();
    }

    @Override
    public int size() {
        return primaryMap.size();
    }

    public Stream>> getRemainingItems() {
        return primaryMap.entrySet().stream().map(kvp -> Map.entry(kvp.getKey(), kvp.getValue().future));
    }
}