package com.netflix.curator.test;

import org.apache.zookeeper.server.ServerCnxnFactory;
import org.apache.zookeeper.server.ServerConfig;
import org.apache.zookeeper.server.ZooKeeperServerMain;
import org.apache.zookeeper.server.quorum.QuorumPeerConfig;
import java.io.IOException;
import java.lang.reflect.Field;
import java.nio.channels.ServerSocketChannel;
import java.util.concurrent.CountDownLatch;

public class TestingZooKeeperMain extends ZooKeeperServerMain implements ZooKeeperMainFace
{
    private final CountDownLatch        latch = new CountDownLatch(1);

    @Override
    public void kill()
    {
        try
        {
            Field               cnxnFactoryField = ZooKeeperServerMain.class.getDeclaredField("cnxnFactory");
            cnxnFactoryField.setAccessible(true);
            ServerCnxnFactory   cnxnFactory = (ServerCnxnFactory)cnxnFactoryField.get(this);
            cnxnFactory.closeAll();

            Field               ssField = cnxnFactory.getClass().getDeclaredField("ss");
            ssField.setAccessible(true);
            ServerSocketChannel ss = (ServerSocketChannel)ssField.get(cnxnFactory);
            ss.close();

            close();
        }
        catch ( Exception e )
        {
            e.printStackTrace();
        }
    }

    @Override
    public void runFromConfig(QuorumPeerConfig config) throws Exception
    {
        ServerConfig        serverConfig = new ServerConfig();
        serverConfig.readFrom(config);
        latch.countDown();
        super.runFromConfig(serverConfig);
    }

    @Override
    public void blockUntilStarted() throws Exception
    {
        latch.await();
        Thread.sleep(100);
    }

    @Override
    public void close() throws IOException
    {
        shutdown();
    }
}

@@ -127,24 +127,123 @@
 rverConfig;%0A
+import org.apache.zookeeper.server.ZKDatabase;%0Aimport org.apache.zookeeper.server.ZooKeeperServer;%0A
 import org.a
@@ -1856,21 +1856,926 @@
  shutdown();
+%0A%0A        try%0A        %7B%0A            Field               cnxnFactoryField = ZooKeeperServerMain.class.getDeclaredField(%22cnxnFactory%22);%0A            cnxnFactoryField.setAccessible(true);%0A            ServerCnxnFactory   cnxnFactory = (ServerCnxnFactory)cnxnFactoryField.get(this);%0A%0A            Field               zkServerField = ServerCnxnFactory.class.getDeclaredField(%22zkServer%22);%0A            zkServerField.setAccessible(true);%0A            ZooKeeperServer     zkServer = (ZooKeeperServer)zkServerField.get(cnxnFactory);%0A            if ( zkServer != null ) %0A            %7B%0A                ZKDatabase      zkDb = zkServer.getZKDatabase();%0A                if ( zkDb != null )%0A                %7B%0A                    // make ZK server close its log files%0A                    zkDb.close();%0A                %7D%0A            %7D%0A        %7D%0A        catch ( Exception e )%0A        %7B%0A            e.printStackTrace();%0A        %7D
 %0A    %7D%0A%7D%0A

Add a hack to stop ZK server leaking file handles on shutdown.