<commit_before>package org.neo4j.kernel.ha;

import java.util.HashMap;
import java.util.LinkedList;
import java.util.Map;
import java.util.Queue;

import org.neo4j.kernel.CommonFactories;
import org.neo4j.kernel.IdGeneratorFactory;
import org.neo4j.kernel.IdType;
import org.neo4j.kernel.ha.zookeeper.ZooKeeperException;
import org.neo4j.kernel.impl.ha.Broker;
import org.neo4j.kernel.impl.ha.IdAllocation;
import org.neo4j.kernel.impl.ha.ResponseReceiver;
import org.neo4j.kernel.impl.nioneo.store.IdGenerator;
import org.neo4j.kernel.impl.nioneo.store.NeoStore;

public class SlaveIdGenerator implements IdGenerator
{
    public static class SlaveIdGeneratorFactory implements IdGeneratorFactory
    {
        private final Broker broker;
        private final ResponseReceiver receiver;
        private final Map<IdType, IdGenerator> generators = new HashMap<IdType, IdGenerator>();
        private final IdGeneratorFactory localFactory =
                CommonFactories.defaultIdGeneratorFactory();

        public SlaveIdGeneratorFactory( Broker broker, ResponseReceiver receiver )
        {
            this.broker = broker;
            this.receiver = receiver;
        }
        
        public IdGenerator open( String fileName, int grabSize, IdType idType, long highestIdInUse )
        {
            IdGenerator localIdGenerator = localFactory.open( fileName, grabSize,
                    idType, highestIdInUse );
            IdGenerator generator = new SlaveIdGenerator( idType, highestIdInUse, broker, receiver,
                    localIdGenerator );
            generators.put( idType, generator );
            return generator;
        }
        
        public void create( String fileName )
        {
            localFactory.create( fileName );
        }

        public IdGenerator get( IdType idType )
        {
            return generators.get( idType );
        }
        
        public void updateIdGenerators( NeoStore store )
        {
            store.updateIdGenerators();
        }
    };
    
    private final Broker broker;
    private final ResponseReceiver receiver;
    private volatile long highestIdInUse;
    private volatile long defragCount;
    private final Queue<Long> idQueue = new LinkedList<Long>();
    private final IdType idType;
    private final IdGenerator localIdGenerator;

    public SlaveIdGenerator( IdType idType, long highestIdInUse, Broker broker,
            ResponseReceiver receiver, IdGenerator localIdGenerator )
    {
        this.idType = idType;
        this.highestIdInUse = highestIdInUse;
        this.broker = broker;
        this.receiver = receiver;
        this.localIdGenerator = localIdGenerator;
    }

    public void close()
    {
        this.localIdGenerator.close();
    }

    public void freeId( long id )
    {
    }

    public long getHighId()
    {
        return this.highestIdInUse;
    }

    public long getNumberOfIdsInUse()
    {
        return this.highestIdInUse - this.defragCount;
    }

    public synchronized long nextId()
    {
        // if we dont have anymore grabbed ids from master, grab a bunch 
        try
        {
            Long nextId = nextLocalIdOrNull();
            if ( nextId == null )
            {
                IdAllocation allocation = receiver.receive(
                        broker.getMaster().allocateIds( receiver.getSlaveContext(), idType ) );
                nextId = storeLocally( allocation );
            }
            return nextId.intValue();
        }
        catch ( ZooKeeperException e )
        {
            receiver.somethingIsWrong( e );
            throw e;
        }
        catch ( HaCommunicationException e )
        {
            receiver.somethingIsWrong( e );
            throw e;
        }
    }

    private Long storeLocally( IdAllocation allocation )
    {
        this.highestIdInUse = allocation.getHighestIdInUse();
        this.defragCount = allocation.getDefragCount();
        for ( long id : allocation.getIds() )
        {
            idQueue.add( id );
        }
        updateLocalIdGenerator();
        return idQueue.poll();
    }

    private void updateLocalIdGenerator()
    {
        long localHighId = this.localIdGenerator.getHighId();
        if ( this.highestIdInUse > localHighId )
        {
            this.localIdGenerator.setHighId( this.highestIdInUse );
        }
    }

    private Long nextLocalIdOrNull()
    {
        return this.idQueue.poll();
    }

    public void setHighId( long id )
    {
        this.highestIdInUse = id;
        this.localIdGenerator.setHighId( id );
    }
    
    public long getDefragCount()
    {
        return this.defragCount;
    }
}<commit_msg>Optimized the storing of ids locally<commit_after>package org.neo4j.kernel.ha;

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

import org.neo4j.kernel.CommonFactories;
import org.neo4j.kernel.IdGeneratorFactory;
import org.neo4j.kernel.IdType;
import org.neo4j.kernel.ha.zookeeper.ZooKeeperException;
import org.neo4j.kernel.impl.ha.Broker;
import org.neo4j.kernel.impl.ha.IdAllocation;
import org.neo4j.kernel.impl.ha.ResponseReceiver;
import org.neo4j.kernel.impl.nioneo.store.IdGenerator;
import org.neo4j.kernel.impl.nioneo.store.NeoStore;

public class SlaveIdGenerator implements IdGenerator
{
    private static final long VALUE_REPRESENTING_NULL = -1;
    
    public static class SlaveIdGeneratorFactory implements IdGeneratorFactory
    {
        private final Broker broker;
        private final ResponseReceiver receiver;
        private final Map<IdType, IdGenerator> generators = new HashMap<IdType, IdGenerator>();
        private final IdGeneratorFactory localFactory =
                CommonFactories.defaultIdGeneratorFactory();

        public SlaveIdGeneratorFactory( Broker broker, ResponseReceiver receiver )
        {
            this.broker = broker;
            this.receiver = receiver;
        }
        
        public IdGenerator open( String fileName, int grabSize, IdType idType, long highestIdInUse )
        {
            IdGenerator localIdGenerator = localFactory.open( fileName, grabSize,
                    idType, highestIdInUse );
            IdGenerator generator = new SlaveIdGenerator( idType, highestIdInUse, broker, receiver,
                    localIdGenerator );
            generators.put( idType, generator );
            return generator;
        }
        
        public void create( String fileName )
        {
            localFactory.create( fileName );
        }

        public IdGenerator get( IdType idType )
        {
            return generators.get( idType );
        }
        
        public void updateIdGenerators( NeoStore store )
        {
            store.updateIdGenerators();
        }
    };
    
    private final Broker broker;
    private final ResponseReceiver receiver;
    private volatile long highestIdInUse;
    private volatile long defragCount;
    private LongArrayIterator idQueue = new LongArrayIterator( new long[0] );
    private final IdType idType;
    private final IdGenerator localIdGenerator;

    public SlaveIdGenerator( IdType idType, long highestIdInUse, Broker broker,
            ResponseReceiver receiver, IdGenerator localIdGenerator )
    {
        this.idType = idType;
        this.highestIdInUse = highestIdInUse;
        this.broker = broker;
        this.receiver = receiver;
        this.localIdGenerator = localIdGenerator;
    }

    public void close()
    {
        this.localIdGenerator.close();
    }

    public void freeId( long id )
    {
    }

    public long getHighId()
    {
        return this.highestIdInUse;
    }

    public long getNumberOfIdsInUse()
    {
        return this.highestIdInUse - this.defragCount;
    }

    public synchronized long nextId()
    {
        try
        {
            long nextId = nextLocalId();
            if ( nextId == VALUE_REPRESENTING_NULL )
            {
                // If we dont have anymore grabbed ids from master, grab a bunch 
                IdAllocation allocation = receiver.receive(
                        broker.getMaster().allocateIds( receiver.getSlaveContext(), idType ) );
                nextId = storeLocally( allocation );
            }
            return nextId;
        }
        catch ( ZooKeeperException e )
        {
            receiver.somethingIsWrong( e );
            throw e;
        }
        catch ( HaCommunicationException e )
        {
            receiver.somethingIsWrong( e );
            throw e;
        }
    }

    private long storeLocally( IdAllocation allocation )
    {
        this.highestIdInUse = allocation.getHighestIdInUse();
        this.defragCount = allocation.getDefragCount();
        this.idQueue = new LongArrayIterator( allocation.getIds() );
        updateLocalIdGenerator();
        return idQueue.next();
    }

    private void updateLocalIdGenerator()
    {
        long localHighId = this.localIdGenerator.getHighId();
        if ( this.highestIdInUse > localHighId )
        {
            this.localIdGenerator.setHighId( this.highestIdInUse );
        }
    }

    private long nextLocalId()
    {
        return this.idQueue.next();
    }

    public void setHighId( long id )
    {
        this.highestIdInUse = id;
        this.localIdGenerator.setHighId( id );
    }
    
    public long getDefragCount()
    {
        return this.defragCount;
    }
    
    private static class LongArrayIterator
    {
        private int position;
        private final long[] array;
        
        LongArrayIterator( long[] array )
        {
            this.array = array;
        }
        
        long next()
        {
            return position < array.length ? array[position++] : VALUE_REPRESENTING_NULL;
        }
    }
}
{'fraction_non_alphanumeric': 0.07615724109617161, 'fraction_numerical': 0.0022580314071641177, 'mean_word_length': 2.0982511923688394, 'pattern_counts': {'":': 0, '<': 10, '<?xml version=': 0, '>': 11, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7053719', 'n_tokens_mistral': 2683, 'n_tokens_neox': 2606, 'n_words': 611}