apoc.util.collection.AbstractResourceIterable Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of apoc-common Show documentation
Data types package for Neo4j Procedures
There is a newer version: 5.25.1
package apoc.util.collection;

import java.util.Arrays;
import java.util.BitSet;
import java.util.Objects;
import java.util.function.Consumer;
import java.util.function.Function;
import org.neo4j.graphdb.ResourceIterable;
import org.neo4j.graphdb.ResourceIterator;

public abstract class AbstractResourceIterable implements ResourceIterable
{
    // start with 2 as most cases will only generate a single iterator but gives a little leeway to save on expansions
    private TrackingResourceIterator[] trackedIterators = new TrackingResourceIterator[2];

    private BitSet trackedIteratorsInUse = new BitSet();

    private boolean closed;

    protected abstract ResourceIterator newIterator();

    @Override
    public final ResourceIterator iterator()
    {
        if ( closed )
        {
            throw new ResourceIteratorCloseFailedException( ResourceIterable.class.getSimpleName() + " has already been closed" );
        }

        return new TrackingResourceIterator<>( Objects.requireNonNull( newIterator() ), this::register, this::unregister );
    }

    @Override
    public final void close()
    {
        if ( !closed )
        {
            try
            {
                internalClose();
            }
            finally
            {
                closed = true;
                onClosed();
            }
        }
    }

    /**
     * Callback method that allows subclasses to perform their own specific closing logic
     */
    protected void onClosed()
    {
    }

    private void register( TrackingResourceIterator iterator )
    {
        if ( trackedIteratorsInUse.cardinality() == trackedIterators.length )
        {
            trackedIterators = Arrays.copyOf( trackedIterators, trackedIterators.length << 1 );
        }

        final var freeIndex = trackedIteratorsInUse.nextClearBit( 0 );
        trackedIterators[freeIndex] = iterator;
        trackedIteratorsInUse.set( freeIndex );
    }

    private void unregister( TrackingResourceIterator iterator )
    {
        final var lastSetBit = trackedIteratorsInUse.previousSetBit( trackedIterators.length );
        for ( int i = 0; i <= lastSetBit; i++ )
        {
            if ( trackedIterators[i] == iterator )
            {
                trackedIterators[i] = null;
                trackedIteratorsInUse.clear( i );
                break;
            }
        }
    }

    private void internalClose()
    {
        ResourceIteratorCloseFailedException closeThrowable = null;
        final var lastSetBit = trackedIteratorsInUse.previousSetBit( trackedIterators.length );
        for ( int i = 0; i <= lastSetBit; i++ )
        {
            if ( trackedIterators[i] == null )
            {
                continue;
            }

            try
            {
                trackedIterators[i].internalClose();
            }
            catch ( Exception e )
            {
                if ( closeThrowable == null )
                {
                    closeThrowable = new ResourceIteratorCloseFailedException( "Exception closing a resource iterator.", e );
                }
                else
                {
                    closeThrowable.addSuppressed( e );
                }
            }
        }

        trackedIterators = null;
        trackedIteratorsInUse = null;

        if ( closeThrowable != null )
        {
            throw closeThrowable;
        }
    }

    private static final class TrackingResourceIterator implements ResourceIterator
    {
        private final ResourceIterator delegate;
        private final Consumer> registerCallback;
        private final Consumer> unregisterCallback;

        private boolean closed;

        private TrackingResourceIterator( ResourceIterator delegate, Consumer> registerCallback,
                Consumer> unregisterCallback )
        {
            this.delegate = delegate;
            this.registerCallback = registerCallback;
            this.unregisterCallback = unregisterCallback;

            registerCallback.accept( this );
        }

        @Override
        public boolean hasNext()
        {
            boolean hasNext = delegate.hasNext();
            if ( !hasNext )
            {
                close();
            }
            return hasNext;
        }

        @Override
        public T next()
        {
            return delegate.next();
        }

        @Override
        public  ResourceIterator map( Function map )
        {
            return new TrackingResourceIterator<>( ResourceIterator.super.map( map ), registerCallback, unregisterCallback );
        }

        @Override
        public void close()
        {
            if ( !closed )
            {
                internalClose();
                unregisterCallback.accept( this );
            }
        }

        private void internalClose()
        {
            try
            {
                delegate.close();
            }
            finally
            {
                closed = true;
            }
        }
    }
}