Repository: ignite
Updated Branches:
  refs/heads/ignite-2.5 c70d85aa3 -> 8d3755b9c


http://git-wip-us.apache.org/repos/asf/ignite/blob/8d3755b9/modules/zookeeper/src/test/java/org/apache/zookeeper/ZkTestClientCnxnSocketNIO.java
----------------------------------------------------------------------
diff --git 
a/modules/zookeeper/src/test/java/org/apache/zookeeper/ZkTestClientCnxnSocketNIO.java
 
b/modules/zookeeper/src/test/java/org/apache/zookeeper/ZkTestClientCnxnSocketNIO.java
new file mode 100644
index 0000000..7892b5e
--- /dev/null
+++ 
b/modules/zookeeper/src/test/java/org/apache/zookeeper/ZkTestClientCnxnSocketNIO.java
@@ -0,0 +1,137 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License.  You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.zookeeper;
+
+import java.io.IOException;
+import java.net.InetSocketAddress;
+import java.nio.channels.SelectionKey;
+import java.util.concurrent.ConcurrentHashMap;
+import java.util.concurrent.CountDownLatch;
+import java.util.concurrent.TimeUnit;
+import org.apache.ignite.Ignite;
+import org.apache.ignite.IgniteLogger;
+import org.apache.ignite.logger.java.JavaLogger;
+import org.apache.ignite.testframework.GridTestUtils;
+
+/**
+ *
+ */
+public class ZkTestClientCnxnSocketNIO extends ClientCnxnSocketNIO {
+    /** */
+    public static final IgniteLogger log = new 
JavaLogger().getLogger(ZkTestClientCnxnSocketNIO.class);
+
+    /** */
+    public static volatile boolean DEBUG = false;
+
+    /** */
+    public volatile CountDownLatch blockConnectLatch;
+
+    /** */
+    public static ConcurrentHashMap<String, ZkTestClientCnxnSocketNIO> clients 
= new ConcurrentHashMap<>();
+
+    /** */
+    private final String nodeName;
+
+    /**
+     *
+     */
+    public static void reset() {
+        clients.clear();
+    }
+
+    /**
+     * @param node Node.
+     * @return ZK client.
+     */
+    public static ZkTestClientCnxnSocketNIO forNode(Ignite node) {
+        return clients.get(node.name());
+    }
+
+    /**
+     * @param instanceName Ignite instance name.
+     * @return ZK client.
+     */
+    public static ZkTestClientCnxnSocketNIO forNode(String instanceName) {
+        return clients.get(instanceName);
+    }
+
+    /**
+     * @throws IOException If failed.
+     */
+    public ZkTestClientCnxnSocketNIO() throws IOException {
+        super();
+
+        String threadName = Thread.currentThread().getName();
+
+        nodeName = threadName.substring(threadName.indexOf('-') + 1);
+
+        if (DEBUG)
+            log.info("ZkTestClientCnxnSocketNIO created for node: " + 
nodeName);
+    }
+
+    /** {@inheritDoc} */
+    @Override void connect(InetSocketAddress addr) throws IOException {
+        CountDownLatch blockConnect = this.blockConnectLatch;
+
+        if (DEBUG)
+            log.info("ZkTestClientCnxnSocketNIO connect [node=" + nodeName + 
", addr=" + addr + ']');
+
+        if (blockConnect != null && blockConnect.getCount() > 0) {
+            try {
+                log.info("ZkTestClientCnxnSocketNIO block connect");
+
+                blockConnect.await(60, TimeUnit.SECONDS);
+
+                log.info("ZkTestClientCnxnSocketNIO finish block connect");
+            }
+            catch (Exception e) {
+                log.error("Error in ZkTestClientCnxnSocketNIO: " + e, e);
+            }
+        }
+
+        super.connect(addr);
+
+        clients.put(nodeName, this);
+    }
+
+    /**
+     *
+     */
+    public void allowConnect() {
+        assert blockConnectLatch != null && blockConnectLatch.getCount() == 1 
: blockConnectLatch;
+
+        log.info("ZkTestClientCnxnSocketNIO allowConnect [node=" + nodeName + 
']');
+
+        blockConnectLatch.countDown();
+    }
+
+    /**
+     * @param blockConnect {@code True} to block client reconnect.
+     * @throws Exception If failed.
+     */
+    public void closeSocket(boolean blockConnect) throws Exception {
+        if (blockConnect)
+            blockConnectLatch = new CountDownLatch(1);
+
+        log.info("ZkTestClientCnxnSocketNIO closeSocket [node=" + nodeName + 
", block=" + blockConnect + ']');
+
+        SelectionKey k = GridTestUtils.getFieldValue(this, 
ClientCnxnSocketNIO.class, "sockKey");
+
+        k.channel().close();
+    }
+}

Reply via email to