added test for RR policy, more efficient counter

hector-client · Dec 14, 2010 · ac02a4c · ac02a4c
1 parent 26a7824
commit ac02a4c
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 35 deletions.
diff --git a/core/src/main/java/me/prettyprint/cassandra/connection/RoundRobinBalancingPolicy.java b/core/src/main/java/me/prettyprint/cassandra/connection/RoundRobinBalancingPolicy.java
@@ -2,18 +2,18 @@
 
 import java.util.Collection;
 import java.util.Set;
+import java.util.concurrent.atomic.AtomicLong;
 
 import org.cliffc.high_scale_lib.Counter;
 
 import me.prettyprint.cassandra.service.CassandraHost;
 
 public class RoundRobinBalancingPolicy implements LoadBalancingPolicy {
 
-  private final Counter counter;
+  private AtomicLong counter;
 
   public RoundRobinBalancingPolicy() {
-    counter = new Counter();
-    counter.set(0);
+    counter = new AtomicLong(Long.MIN_VALUE);
   }
 
   @Override
@@ -24,16 +24,14 @@ public ConcurrentHClientPool getPool(Collection<ConcurrentHClientPool> pools,
     ConcurrentHClientPool pool = pa[location];
     if ( excludeHosts != null && excludeHosts.size() > 0 ) {
       while ( excludeHosts.contains(pool.getCassandraHost()) ) {
-        pool = pa[++location == pa.length ? 0 : location]; 
+        pool = pa[getAndIncrement(pa.length)]; 
       }
     }
     return pool;
   }
-
+    
   private int getAndIncrement(int size) {
-    int val = (int)(counter.get() % size); 
-    counter.increment();
-    return val;
+    return (int)counter.getAndIncrement() % size;    
   }
 
 }
diff --git a/core/src/test/java/me/prettyprint/cassandra/connection/LeastActiveBalancingPolicyTest.java b/core/src/test/java/me/prettyprint/cassandra/connection/LeastActiveBalancingPolicyTest.java
@@ -1,42 +1,18 @@
 package me.prettyprint.cassandra.connection;
 
-import static org.junit.Assert.*;
+import static org.junit.Assert.assertEquals;
 
-import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.HashSet;
-import java.util.List;
 
 import me.prettyprint.cassandra.service.CassandraHost;
 
-import org.junit.Before;
 import org.junit.Test;
-import org.mockito.Mock;
 import org.mockito.Mockito;
 
-public class LeastActiveBalancingPolicyTest {
+public class LeastActiveBalancingPolicyTest extends BaseBalancingPolicyTest {
 
   private LeastActiveBalancingPolicy leastActiveBalancingPolicy;
-  private List<ConcurrentHClientPool> pools = new ArrayList<ConcurrentHClientPool>();
-
-  private ConcurrentHClientPool poolWith5Active;
-  private ConcurrentHClientPool poolWith7Active;
-  private ConcurrentHClientPool poolWith10Active;
-
-
-  @Before
-  public void setup() {
-    poolWith5Active = Mockito.mock(ConcurrentHClientPool.class);    
-    Mockito.when(poolWith5Active.getNumActive()).thenReturn(5);
-    poolWith7Active = Mockito.mock(ConcurrentHClientPool.class);
-    Mockito.when(poolWith7Active.getNumActive()).thenReturn(7);
-    poolWith10Active = Mockito.mock(ConcurrentHClientPool.class);
-    Mockito.when(poolWith10Active.getNumActive()).thenReturn(10);
-
-    pools.add(poolWith5Active);
-    pools.add(poolWith7Active);
-    pools.add(poolWith10Active);
-  }
 
   @Test
   public void testGetPoolOk() {

diff --git a/core/src/test/java/me/prettyprint/cassandra/connection/RoundRobinBalancingPolicyTest.java b/core/src/test/java/me/prettyprint/cassandra/connection/RoundRobinBalancingPolicyTest.java
@@ -2,7 +2,13 @@
 
 import static org.junit.Assert.*;
 
+import java.util.Arrays;
+import java.util.HashSet;
+
+import me.prettyprint.cassandra.service.CassandraHost;
+
 import org.junit.Test;
+import org.mockito.Mockito;
 
 
 public class RoundRobinBalancingPolicyTest extends BaseBalancingPolicyTest {
@@ -20,7 +26,20 @@ public void testGetPoolOk() {
     assertEquals(poolWith10Active, roundRobinBalancingPolicy.getPool(pools, null));
     assertEquals(poolWith5Active, roundRobinBalancingPolicy.getPool(pools, null));
     assertEquals(poolWith7Active, roundRobinBalancingPolicy.getPool(pools, null));
-    assertEquals(poolWith10Active, roundRobinBalancingPolicy.getPool(pools, null));
+    assertEquals(poolWith10Active, roundRobinBalancingPolicy.getPool(pools, null));    
+  }
+
+  @Test
+  public void testSkipExhausted() {
+    Mockito.when(poolWith5Active.getCassandraHost()).thenReturn(new CassandraHost("127.0.0.1:9160"));
+    Mockito.when(poolWith7Active.getCassandraHost()).thenReturn(new CassandraHost("127.0.0.2:9161"));
+    Mockito.when(poolWith10Active.getCassandraHost()).thenReturn(new CassandraHost("127.0.0.3:9162"));
+
+    roundRobinBalancingPolicy = new RoundRobinBalancingPolicy();
+    assertEquals(poolWith7Active, roundRobinBalancingPolicy.getPool(pools, new HashSet<CassandraHost>(Arrays.asList(new CassandraHost("127.0.0.1:9160")))));
+    assertEquals(poolWith10Active, roundRobinBalancingPolicy.getPool(pools, new HashSet<CassandraHost>(Arrays.asList(new CassandraHost("127.0.0.1:9160")))));
+    assertEquals(poolWith7Active, roundRobinBalancingPolicy.getPool(pools, new HashSet<CassandraHost>(Arrays.asList(new CassandraHost("127.0.0.1:9160")))));
+    assertEquals(poolWith10Active, roundRobinBalancingPolicy.getPool(pools, new HashSet<CassandraHost>(Arrays.asList(new CassandraHost("127.0.0.1:9160")))));
   }
 
 }