Remove duplicate start entries from logical log

package org.neo4j;

import static java.nio.ByteBuffer.allocate;
import static org.neo4j.kernel.impl.transaction.xaframework.LogIoUtils.readEntry;
import static org.neo4j.kernel.impl.transaction.xaframework.LogIoUtils.readLogHeader;
import static org.neo4j.kernel.impl.transaction.xaframework.LogIoUtils.writeLogEntry;
import static org.neo4j.kernel.impl.transaction.xaframework.LogIoUtils.writeLogHeader;

import java.io.File;
import java.io.IOException;
import java.io.RandomAccessFile;
import java.lang.reflect.Method;
import java.nio.ByteBuffer;
import java.nio.channels.FileChannel;
import java.nio.channels.ReadableByteChannel;
import java.util.HashSet;
import java.util.Set;

import javax.transaction.xa.Xid;

import org.neo4j.index.impl.lucene.LuceneDataSource;
import org.neo4j.kernel.impl.nioneo.xa.Command;
import org.neo4j.kernel.impl.nioneo.xa.NeoStoreXaDataSource;
import org.neo4j.kernel.impl.transaction.xaframework.DirectMappedLogBuffer;
import org.neo4j.kernel.impl.transaction.xaframework.LogBuffer;
import org.neo4j.kernel.impl.transaction.xaframework.LogEntry;
import org.neo4j.kernel.impl.transaction.xaframework.XaCommand;
import org.neo4j.kernel.impl.transaction.xaframework.XaCommandFactory;

public class RemoveDualStartRecords
{
    public static void main( String[] args ) throws Exception
    {
        if ( args.length < 1 ) exit( "Supply store directory to fix" );
        File dir = new File( args[0] );
        if ( !dir.exists() ) exit( "Invalid store directory " + dir.getAbsolutePath() );

        CommandFactory neoCommands = new CommandFactory();
        fix( new File( dir, NeoStoreXaDataSource.LOGICAL_LOG_DEFAULT_NAME + ".1" ), neoCommands );
        fix( new File( dir, NeoStoreXaDataSource.LOGICAL_LOG_DEFAULT_NAME + ".2" ), neoCommands );

        IndexCommandFactory indexCommands = new IndexCommandFactory();
        fix( new File( new File( dir, "index" ), "lucene.log.1" ), indexCommands );
        fix( new File( new File( dir, "index" ), "lucene.log.2" ), indexCommands );
    }

    private static void fix( File file, XaCommandFactory cf ) throws IOException
    {
        if ( !file.exists() ) return;
        FileChannel inChannel = new RandomAccessFile( file, "r" ).getChannel();
        ByteBuffer buffer = allocate( 500 );
        long[] header = readLogHeader( buffer, inChannel, true );
        File fixedFile = new File( file.getAbsolutePath() + ".fixed" );
        writeLogHeader( buffer, header[0], header[1] );
        FileChannel outChannel = new RandomAccessFile( fixedFile, "rw" ).getChannel();
        outChannel.write( buffer ); // header
        LogBuffer logBuffer = new DirectMappedLogBuffer( outChannel );
        LogEntry entry = null;
        Set<Xid> startXids = new HashSet<Xid>();
        boolean success = false;
        int removedCount = 0;
        try
        {
            while ( (entry = readEntry( buffer, inChannel, cf )) != null )
            {
                if ( entry instanceof LogEntry.Start && !startXids.add( ((LogEntry.Start) entry).getXid() ) )
                {
                    removedCount++;
                    continue;
                }
                writeLogEntry( entry, logBuffer );
            }
            success = true;
        }
        finally
        {
            logBuffer.force();
            inChannel.close();
            outChannel.close();
        }

        if ( success )
        {
            file.delete();
            fixedFile.renameTo( file );
            if ( removedCount != 0 ) System.out.println( "Removed " + removedCount + " duplicate entries from " + file.getAbsolutePath() );
        }
    }

    private static void exit( String string )
    {
        System.out.println( string );
        System.exit( 1 );
    }

    private static class CommandFactory extends XaCommandFactory
    {
        @Override
        public XaCommand readCommand( ReadableByteChannel byteChannel,
                ByteBuffer buffer ) throws IOException
        {
            return Command.readCommand( null, byteChannel, buffer );
        }
    }

    private static class IndexCommandFactory extends XaCommandFactory
    {
        private final Class<?> luceneCommandClass;
        private final Method readCommandMethod;

        IndexCommandFactory() throws Exception
        {
            luceneCommandClass = Class.forName( "org.neo4j.index.impl.lucene.LuceneCommand" );
            readCommandMethod = luceneCommandClass.getDeclaredMethod( "readCommand", ReadableByteChannel.class, ByteBuffer.class, LuceneDataSource.class );
            readCommandMethod.setAccessible( true );
        }

        @Override
        public XaCommand readCommand( ReadableByteChannel byteChannel, ByteBuffer buffer )
                throws IOException
        {
//            return LuceneCommand.readCommand( byteChannel, buffer, null );
            try
            {
                return (XaCommand) readCommandMethod.invoke( null, byteChannel, buffer, null );
            }
            catch ( Exception e )
            {
                throw new RuntimeException( e );
            }
        }
    }
}