package com.datastax.driver.core;

import java.util.Iterator;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicInteger;

import org.testng.annotations.Test;

import static org.assertj.core.api.Assertions.assertThat;

import com.datastax.driver.core.policies.DelegatingLoadBalancingPolicy;
import com.datastax.driver.core.policies.LoadBalancingPolicy;
import com.datastax.driver.core.policies.Policies;
import com.datastax.driver.core.policies.ReconnectionPolicy;

public class ControlConnectionTest {
    @Test(groups = "short")
    public void should_prevent_simultaneous_reconnection_attempts() throws InterruptedException {
        CCMBridge ccm = null;
        Cluster cluster = null;

        // Custom load balancing policy that counts the number of calls to newQueryPlan().
        // Since we don't open any session from our Cluster, only the control connection reattempts are calling this
        // method, therefore the invocation count is equal to the number of attempts.
        QueryPlanCountingPolicy loadBalancingPolicy = new QueryPlanCountingPolicy(Policies.defaultLoadBalancingPolicy());

        // Custom reconnection policy with a very large delay (longer than the test duration), to make sure we count
        // only the first reconnection attempt of each reconnection handler.
        ReconnectionPolicy reconnectionPolicy = new ReconnectionPolicy() {
            @Override
            public ReconnectionSchedule newSchedule() {
                return new ReconnectionSchedule() {
                    @Override
                    public long nextDelayMs() {
                        return 60 * 1000;
                    }
                };
            }
        };

        try {
            ccm = CCMBridge.create("test", 1);
            cluster = Cluster.builder()
                             .addContactPoint(CCMBridge.ipOfNode(1))
                             .withReconnectionPolicy(reconnectionPolicy)
                             .withLoadBalancingPolicy(loadBalancingPolicy)
                             .build();
            cluster.init();

            ccm.stop(1);

            // Sleep for a while to make sure our final count is not the result of lucky timing
            TimeUnit.SECONDS.sleep(1);

        } finally {
            if (cluster != null)
                cluster.close();
            if (ccm != null)
                ccm.remove();
        }
        assertThat(loadBalancingPolicy.counter.get()).isEqualTo(1);
    }

    static class QueryPlanCountingPolicy extends DelegatingLoadBalancingPolicy {

        final AtomicInteger counter = new AtomicInteger();

        public QueryPlanCountingPolicy(LoadBalancingPolicy delegate) {
            super(delegate);
        }

        public Iterator<Host> newQueryPlan(String loggedKeyspace, Statement statement) {
            counter.incrementAndGet();
            return super.newQueryPlan(loggedKeyspace, statement);
        }
    }
}

@@ -1128,16 +1128,90 @@
 licy());
+%0A        AtomicInteger reconnectionAttempts = loadBalancingPolicy.counter;
 %0A%0A      
@@ -1847,20 +1847,136 @@
 %22test%22, 
-1
+2
 );%0A
+            // We pass only the first host as contact point, so we know the control connection will be on this host%0A
         
@@ -2007,29 +2007,16 @@
 ilder()%0A
-             
         
@@ -2075,37 +2075,24 @@
             
-          
-   
 .withReconne
@@ -2119,37 +2119,24 @@
 tionPolicy)%0A
-             
             
@@ -2193,32 +2193,42 @@
                 
+.build();%0A
              .bu
@@ -2219,35 +2219,41 @@
             
- .build
+cluster.init
 ();%0A
+%0A
             
@@ -2252,32 +2252,101 @@
         
-cluster.init();%0A
+// Kill the control connection host, there should be exactly one reconnection attempt
 %0A       
@@ -2359,25 +2359,24 @@
 cm.stop(1);%0A
-%0A
             
@@ -2370,24 +2370,51 @@
 %0A           
+ TimeUnit.SECONDS.sleep(1);
  // Sleep fo
@@ -2502,32 +2502,188 @@
     
-TimeUnit.SECONDS.sleep(1
+assertThat(reconnectionAttempts.get()).isEqualTo(1);%0A%0A            ccm.stop(2);%0A            TimeUnit.SECONDS.sleep(1);%0A            assertThat(reconnectionAttempts.get()).isEqualTo(2
 );%0A%0A
@@ -2841,76 +2841,8 @@
   %7D%0A
-        assertThat(loadBalancingPolicy.counter.get()).isEqualTo(1);%0A
     

Improve test for control connection reconnection (JAVA-497).