adjusted endpoint type to number of host

VPCE have a number of endpoints depending on the number of availablity zones for a region.

Author: michaelraney

README.md

@@ -40,7 +40,7 @@ To activate this throttler, modify the {@code advanced.throttler} section in the
          advanced.throttler = {
                 class = com.aws.ssa.keyspaces.throttler.AmazonKeyspacesFixedRateThrottler
                 max-requests-per-second = 1000
-                endpoint-type = VPC
+                number-of-hosts = 3
                 register-timeout = 1 seconds
           }
       }  
@@ -49,7 +49,7 @@ To activate this throttler, modify the {@code advanced.throttler} section in the
 
 
 * `max-requests-per-second` : controls the request rate. Blocks until available permits or timeout it reached
-* `endpoint-type` : if you are connected to the VPC endpoint or the public endpoint. Used to determine throughput based on the number of connections specified in:`advanced.connection.pool.local.size`
+* `number-of-hosts` : The number of hosts in the system.peers table.  Depending on the endpoint type and region the number of hosts in the system.peers table may be different. This number is Used to validate throughput based on the number of connections specified in:`advanced.connection.pool.local.size`
 * `register-timeout` timeout waiting for permits. Should be less than or equal to `basic.request.timeout'
 
 

src/main/java/com/aws/ssa/keyspaces/core/EndpointType.java

@@ -134,7 +134,6 @@ public RetryDecision onWriteTimeout(
 
         LOG.trace(RETRYING_ON_WRITE_TIMEOUT, logPrefix, cl, blockFor, received, false, retryCount);
 
-
         return decision;
     }
 
@@ -157,7 +156,6 @@ public RetryDecision onUnavailable(
 
         LOG.trace(RETRYING_ON_UNAVAILABLE, logPrefix, cl, required, alive, retryCount);
 
-
         return decision;
     }
 
@@ -174,7 +172,6 @@ public RetryDecision onRequestAborted(
 
         LOG.trace(RETRYING_ON_ABORTED, logPrefix, retryCount, error);
 
-
         return decision;
     }
 

src/main/java/com/aws/ssa/keyspaces/retry/AmazonKeyspacesRetryPolicy.java

@@ -1,6 +1,5 @@
 package com.aws.ssa.keyspaces.throttler;
 
-import com.aws.ssa.keyspaces.core.EndpointType;
 import com.datastax.oss.driver.api.core.RequestThrottlingException;
 import com.datastax.oss.driver.api.core.config.DefaultDriverOption;
 import com.datastax.oss.driver.api.core.context.DriverContext;
@@ -69,10 +68,10 @@ public class AmazonKeyspacesFixedRateThrottler implements RequestThrottler {
     public static int PUBLIC_ENDPOINT_DEFAULT_HOST = 9;
 
     /***
-     * Amazon Keyspaces Virtual Private Cloud Endpoint (VPCE) exposes single host to drivers. The default behavior will establish one connection for every conection specified in
+     * Amazon Keyspaces Virtual Private Cloud Endpoint (VPCE) exposes host per availability zone. The default behavior will establish one connection to each peer IP address.
      * the driver.
      */
-    public static int VPC_ENDPOINT_DEFAULT_HOST = 1;
+    public static int VPC_ENDPOINT_DEFAULT_HOST = 2;
 
     /***
      * Rate limiter used to meter the CQL Request Per Second up to maxRequestsPerSecond
@@ -85,9 +84,9 @@ public class AmazonKeyspacesFixedRateThrottler implements RequestThrottler {
     private final RateLimiter maxConnectionsLimiter;
 
     /***
-     * Type of endpoint used which dictate the number of connections generated when creating a new session
+     * Number of hosts available when creating to a new session
      */
-    private final EndpointType endpointType;
+    private final Integer numberOfHosts;
 
     /***
      * Configured Rate of desired throughput
@@ -116,9 +115,9 @@ public AmazonKeyspacesFixedRateThrottler(DriverContext context) {
                             .getDuration(KeyspacesThrottleOption.KEYSPACES_THROTTLE_TIMEOUT,context.getConfig()
                             .getDefaultProfile()
                             .getDuration(DefaultDriverOption.REQUEST_TIMEOUT)).toMillis(),
-                    EndpointType.valueOf(context.getConfig()
+                    context.getConfig()
                             .getDefaultProfile()
-                            .getString(KeyspacesThrottleOption.KEYSPACES_THROTTLE_ENDPOINT_TYPE, KeyspacesThrottleOption.DEFAULT_ENDPOINT_TYPE.toString())),
+                            .getInt(KeyspacesThrottleOption.KEYSPACES_THROTTLE_NUMBER_OF_HOSTS, KeyspacesThrottleOption.DEFAULT_NUMBER_OF_HOSTS),
                     context.getConfig()
                             .getDefaultProfile()
                             .getInt(DefaultDriverOption.CONNECTION_POOL_LOCAL_SIZE));
@@ -127,14 +126,14 @@ public AmazonKeyspacesFixedRateThrottler(DriverContext context) {
 
         /*** Initialization of the Throttler ***/
         public AmazonKeyspacesFixedRateThrottler(DriverContext context, long maxRequestsPerSecond, long registerTimeoutInMs
-                , EndpointType endpointType, int numberOfConnectionsPerHost) {
+                , int numberOfHosts, int numberOfConnectionsPerHost) {
             this.logPrefix = context.getSessionName();
 
             this.maxRequestsPerSecond = maxRequestsPerSecond;
 
             this.registerTimeoutInMs = registerTimeoutInMs;
 
-            this.endpointType = endpointType;
+            this.numberOfHosts = numberOfHosts;
 
             this.numberOfConnectionsPerHost = numberOfConnectionsPerHost;
 
@@ -179,29 +178,29 @@ public AmazonKeyspacesFixedRateThrottler(DriverContext context, long maxRequests
                     registerTimeoutInMs);
 
             LOG.info(
-                    "[{}] Based on Throttler max of {} request per second, the recommended number of connections for VPC Endpoint is: {}, currently {}",
+                    "[{}] Based on Throttler max of {} request per second, the recommended number of connections for number of hosts: {}, currently {}",
                     logPrefix,
                     maxRequestsPerSecond,
-                    simpleConnectionRecommendation(maxRequestsPerSecond, EndpointType.VPC),
+                    simpleConnectionRecommendation(maxRequestsPerSecond, numberOfHosts),
                     numberOfConnectionsPerHost);
 
             LOG.info(
                     "[{}] Based on Throttler max of {} request per second, the recommended number of connections for the Public Endpoint is: {}, currently {}",
                     logPrefix,
                     maxRequestsPerSecond,
-                    simpleConnectionRecommendation(maxRequestsPerSecond, EndpointType.PUBLIC),
+                    simpleConnectionRecommendation(maxRequestsPerSecond, numberOfHosts),
                     numberOfConnectionsPerHost);
 
 
             //The rate for the number of request per second based on the number of connections. Should be greater than maxRequestsPerSecond
-            int maxRequestPerSecondByForConnections = calculateConnectionMaxRequestPerSecond(endpointType, numberOfConnectionsPerHost);
+            int maxRequestPerSecondByForConnections = calculateConnectionMaxRequestPerSecond(numberOfHosts, numberOfConnectionsPerHost);
 
             if(maxRequestPerSecondByForConnections < maxRequestsPerSecond){
                 LOG.warn(
-                        "[{}] Cannot reach Max Request Per Second of {}. Specified Endpoint type of {}, and number of connections {} will provide at most {} request per second. Try increasing advanced.connection.pool.local.size ",
+                        "[{}] Cannot reach Max Request Per Second of {}. Specified number of hosts {}, and number of connections {} will provide at most {} request per second. Try increasing advanced.connection.pool.local.size or check system.peers table fo the number of hosts ",
                         logPrefix,
                         maxRequestsPerSecond,
-                        endpointType,
+                        numberOfHosts,
                         numberOfConnectionsPerHost,
                         maxRequestPerSecondByForConnections);
 
@@ -218,27 +217,23 @@ public AmazonKeyspacesFixedRateThrottler(DriverContext context, long maxRequests
 
     /***
      * Calculate the number of request per second based on the number of connections, number of hosts, and 2000 request per second.
-     * @param endpointType VPC or Public Endpoint
+     * @param numberOfHosts Number of hosts in the peers table. Depends on region and end point
      * @param numberOfConnectionsPerHost Number of connections for each host ip
      * @return max rate per second based on the number of connections
      */
-        public static int calculateConnectionMaxRequestPerSecond(EndpointType endpointType, int numberOfConnectionsPerHost){
-            return (endpointType == EndpointType.VPC)
-                    ? numberOfConnectionsPerHost * VPC_ENDPOINT_DEFAULT_HOST * REQUEST_PER_CONNECTION_DEFAULT
-                    : numberOfConnectionsPerHost * PUBLIC_ENDPOINT_DEFAULT_HOST * REQUEST_PER_CONNECTION_DEFAULT;
+        public static int calculateConnectionMaxRequestPerSecond(Integer numberOfHosts, int numberOfConnectionsPerHost){
+            return  numberOfConnectionsPerHost * numberOfHosts * REQUEST_PER_CONNECTION_DEFAULT;
         }
 
     /***
      * Calculate recommended connections based on the current maxRequestRate specified. Max Connections rate should be greater than configured max request rate
      * @param maxRequestsPerSecond number of configured request per second
-     * @param endpointType type of endpoint that will be used to identify the number of hosts
+     * @param numberOfHosts type of endpoint that will be used to identify the number of hosts
      * @return
      */
-        public static double simpleConnectionRecommendation(long maxRequestsPerSecond, EndpointType endpointType){
-            int hosts = (endpointType == EndpointType.VPC)
-                    ? VPC_ENDPOINT_DEFAULT_HOST:PUBLIC_ENDPOINT_DEFAULT_HOST;
+        public static double simpleConnectionRecommendation(long maxRequestsPerSecond, Integer numberOfHosts){
 
-            return Math.max(1.0, Math.ceil(maxRequestsPerSecond/(double)(hosts * REQUEST_PER_CONNECTION_DEFAULT)));
+            return Math.max(1.0, Math.ceil(maxRequestsPerSecond/(double)(numberOfHosts * REQUEST_PER_CONNECTION_DEFAULT)));
         }
 
     /***

src/main/java/com/aws/ssa/keyspaces/throttler/AmazonKeyspacesFixedRateThrottler.java

@@ -1,15 +1,13 @@
 package com.aws.ssa.keyspaces.throttler;
 
-import com.aws.ssa.keyspaces.core.EndpointType;
 import com.datastax.oss.driver.api.core.config.DriverOption;
 
 public enum KeyspacesThrottleOption implements DriverOption {
 
-    KEYSPACES_THROTTLE_ENDPOINT_TYPE("advanced.throttler.endpoint-type"),
+    KEYSPACES_THROTTLE_NUMBER_OF_HOSTS("advanced.throttler.number-of-hosts"),
     KEYSPACES_THROTTLE_TIMEOUT("advanced.throttler.register-timeout");;
 
-    public static final EndpointType DEFAULT_ENDPOINT_TYPE = EndpointType.VPC;
-
+    public static final int DEFAULT_NUMBER_OF_HOSTS = 1;
 
     private final String path;
 

src/main/java/com/aws/ssa/keyspaces/throttler/KeyspacesThrottleOption.java

@@ -0,0 +1,145 @@
+package com.aws.ssa.keyspaces.export;
+
+import com.datastax.oss.driver.api.core.ConsistencyLevel;
+import com.datastax.oss.driver.api.core.CqlSession;
+import com.datastax.oss.driver.api.core.config.DriverConfigLoader;
+import com.datastax.oss.driver.api.core.cql.*;
+import com.datastax.oss.driver.shaded.guava.common.util.concurrent.RateLimiter;
+
+import java.nio.ByteBuffer;
+import java.util.*;
+import java.util.concurrent.*;
+import java.util.concurrent.atomic.AtomicLong;
+
+
+public class Exporter {
+
+    private  static BlockingQueue<Map<ByteBuffer,  List<Row>>> blockingQueue = new LinkedBlockingDeque<>(2000);
+
+    private static int PAGE_SIZE = 5000;
+
+    private static String query = "";
+
+    private static ByteBuffer pagingState = null;
+
+    private static AtomicLong count = new AtomicLong(0L);
+
+    private static AtomicLong queueCount = new AtomicLong(0L);
+
+    private RateLimiter limiter = RateLimiter.create(10000);
+
+    public static void main(String[] args){
+
+
+        //query = args[0];
+
+        query = "SELECT * FROM tlp_stress.counter_wide";
+
+        SimpleStatement t = new SimpleStatementBuilder(query).build();
+
+        t.setConsistencyLevel(ConsistencyLevel.LOCAL_ONE);
+        t.setIdempotent(true);
+        t.setPageSize(PAGE_SIZE);
+
+
+        CqlSession session = CqlSession.builder().withConfigLoader(DriverConfigLoader.fromClasspath("throttler-example")).build();
+
+       new Exporter().saveAllRows(session, t);
+
+        ExecutorService executor = Executors.newSingleThreadExecutor();
+        Future task = executor.submit(() -> {
+
+            String threadName = Thread.currentThread().getName();
+            System.out.println("Hello " + threadName);
+
+            do{
+                try {
+                    Map<ByteBuffer, List<Row>> results = blockingQueue.poll(10000, TimeUnit.MILLISECONDS);
+                    //results.entrySet().forEach();
+                    if (results != null) {
+                        results.entrySet().forEach(x -> {
+                            x.getValue().forEach(y -> {
+                                queueCount.incrementAndGet();
+                            });
+                        });
+                        //System.out.println(results);
+                    }
+                }catch (Exception ex){
+
+                }
+                System.out.println("Fin " + count + " " + queueCount);
+            }while(true);
+        });
+
+        while(task.isDone()!=true){
+            try{
+                Thread.sleep(100);
+            }catch (Exception ex){
+
+            }
+        }
+        System.out.println("Fin " + count + " " + queueCount);
+
+    }
+    private void saveAllRows(CqlSession session, SimpleStatement statement){
+
+        CompletionStage<AsyncResultSet> resultSetFuture =
+                session.executeAsync(statement);
+
+        CompletionStage<Void> fetchRows = resultSetFuture.thenCompose(this::queueUpPages);
+
+        fetchRows.whenComplete(
+                (resultSet, error) -> {
+                    if (error != null) {
+                        System.out.printf("Error: %s%n", error.getMessage());
+
+                        SimpleStatement t = new SimpleStatementBuilder(query).build();
+
+                        t.setConsistencyLevel(ConsistencyLevel.LOCAL_ONE);
+                        t.setIdempotent(true);
+                        t.setPageSize(PAGE_SIZE);
+                        t.setPagingState(pagingState);
+
+                        saveAllRows(session, t);
+
+                    } else {
+                        //String version = resultSet.one().getString(0);
+                        System.out.printf("Finished");
+                    }
+                });
+
+
+
+    }
+    private CompletionStage<Void> queueUpPages(AsyncResultSet resultSet)  {
+
+        List<Row> rowsPerPage = new ArrayList<Row>(PAGE_SIZE+50);
+
+        int totalCount = 0;
+
+        for (Row row : resultSet.currentPage()) {
+            count.incrementAndGet();
+            rowsPerPage.add(row);
+            totalCount++;
+        }
+
+        Map<ByteBuffer, List<Row>> oneState = new HashMap<>(1);
+
+        pagingState = resultSet.getExecutionInfo().getPagingState();
+
+        oneState.put(pagingState, rowsPerPage);
+
+        try {
+            blockingQueue.put(oneState);
+        }catch (Exception ex){
+            throw new CompletionException(ex);
+        }
+
+        if (resultSet.hasMorePages()) {
+            limiter.acquire(totalCount);
+            return resultSet.fetchNextPage().thenCompose(this::queueUpPages);
+        } else {
+            return CompletableFuture.completedFuture(null);
+        }
+    }
+}