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

Go to download
package org.opensearch.migrations.replay.util;

import java.util.ArrayDeque;
import java.util.Deque;
import java.util.function.Function;
import java.util.stream.Stream;

import com.google.errorprone.annotations.MustBeClosed;
import io.netty.util.ReferenceCounted;

public final class RefSafeStreamUtils {
    @MustBeClosed
    public static  Stream refSafeMap(
        Stream inputStream,
        Function referenceTrackedMappingFunction
    ) {
        final Deque refCountedTracker = new ArrayDeque<>();
        return inputStream.map(t -> {
            var resource = referenceTrackedMappingFunction.apply(t);
            refCountedTracker.add(resource);
            return resource;
        }).onClose(() -> refCountedTracker.forEach(ReferenceCounted::release));
    }

    public static  U refSafeTransform(
        Stream inputStream,
        Function transformCreatingReferenceTrackedObjects,
        Function, U> streamApplication
    ) {
        try (var mappedStream = refSafeMap(inputStream, transformCreatingReferenceTrackedObjects)) {
            return streamApplication.apply(mappedStream);
        }
    }

    private RefSafeStreamUtils() {}
}