Fix for RetryStats and rate limiting in complex queries.

Fixed problem with rate limiting in complex queries with sorting, where
limits were counted twice resulting in 1/2 of expected throughput.
Also fixed stats not being set properly into QueryResult.

Author: connelly38

driver/src/main/java/oracle/nosql/driver/http/Client.java

@@ -325,6 +325,10 @@ public Result execute(Request kvRequest) {
          */
         kvRequest.validate();
 
+        /* clear any retry stats that may exist on this request object */
+        kvRequest.setRetryStats(null);
+        kvRequest.setRateLimitDelayedMs(0);
+
         if (kvRequest.isQueryRequest()) {
             QueryRequest qreq = (QueryRequest)kvRequest;
 
@@ -377,9 +381,6 @@ public Result execute(Request kvRequest) {
 
         Throwable exception = null;
 
-        /* clear any retry stats that may exist on this request object */
-        kvRequest.setRetryStats(null);
-
         /*
          * If the request doesn't set an explicit compartment, use
          * the config default if provided.
@@ -603,6 +604,12 @@ public Result execute(Request kvRequest) {
                     updateRateLimiters(((TableResult)res).getTableName(), tl);
                 }
 
+                /*
+                 * We may not have rate limiters yet because queries may
+                 * not have a tablename until after the first request.
+                 * So try to get rate limiters if we don't have them yet and
+                 * this is a QueryRequest.
+                 */
                 if (rateLimiterMap != null && readLimiter == null) {
                     readLimiter = getQueryRateLimiter(kvRequest, true);
                 }
@@ -829,9 +836,18 @@ private RateLimiter getQueryRateLimiter(Request request, boolean read) {
         }
 
         if (read) {
-            return rateLimiterMap.getReadLimiter(tableName);
+            RateLimiter rl = rateLimiterMap.getReadLimiter(tableName);
+            if (rl != null) {
+                request.setReadRateLimiter(rl);
+            }
+            return rl;
+        }
+
+        RateLimiter rl = rateLimiterMap.getWriteLimiter(tableName);
+        if (rl != null) {
+            request.setWriteRateLimiter(rl);
         }
-        return rateLimiterMap.getWriteLimiter(tableName);
+        return rl;
     }
 
     /**

driver/src/main/java/oracle/nosql/driver/ops/QueryRequest.java

@@ -114,7 +114,8 @@ public QueryRequest() {
     public QueryRequest copyInternal() {
 
         QueryRequest internalReq = new QueryRequest();
-        internalReq.timeoutMs = timeoutMs;
+        super.copyTo(internalReq);
+
         internalReq.traceLevel = traceLevel;
         internalReq.limit = limit;
         internalReq.maxReadKB = maxReadKB;
@@ -125,7 +126,6 @@ public QueryRequest copyInternal() {
         internalReq.preparedStatement = preparedStatement;
         internalReq.isInternal = true;
         internalReq.driver = driver;
-        internalReq.tableName = tableName;
         return internalReq;
     }
 

driver/src/main/java/oracle/nosql/driver/ops/QueryResult.java

@@ -11,7 +11,6 @@
 
 import oracle.nosql.driver.Consistency;
 import oracle.nosql.driver.NoSQLHandle;
-import oracle.nosql.driver.RateLimiter;
 import oracle.nosql.driver.query.QueryDriver;
 import oracle.nosql.driver.values.MapValue;
 
@@ -230,23 +229,6 @@ void compute() {
         QueryDriver driver = request.getDriver();
         driver.compute(this);
         isComputed = true;
-
-        /*
-         * If the original request specified rate limiting, apply the
-         * used read/write units to the limiter(s) here
-         */
-        if (request != null) {
-            RateLimiter readLimiter = request.getReadRateLimiter();
-            if (readLimiter != null) {
-                readLimiter.consumeUnitsUnconditionally(
-                    super.getReadUnitsInternal());
-            }
-            RateLimiter writeLimiter = request.getWriteRateLimiter();
-            if (writeLimiter != null) {
-                writeLimiter.consumeUnitsUnconditionally(
-                    super.getWriteUnitsInternal());
-            }
-        }
     }
 
     /**

driver/src/main/java/oracle/nosql/driver/ops/Request.java

@@ -423,4 +423,20 @@ public int getRateLimitDelayedMs() {
      * @return the type name of the request
      */
     public abstract String getTypeName();
+
+    /**
+     * @hidden
+     * Copy internal fields to another Request object.
+     */
+    public void copyTo(Request other) {
+        other.setTimeoutInternal(this.timeoutMs);
+        other.setCheckRequestSize(this.checkRequestSize);
+        other.setCompartmentInternal(this.compartment);
+        other.setTableNameInternal(this.tableName);
+        other.setStartTimeMs(this.startTimeMs);
+        other.setRetryStats(this.retryStats);
+        other.setReadRateLimiter(this.readRateLimiter);
+        other.setWriteRateLimiter(this.writeRateLimiter);
+        other.setRateLimitDelayedMs(this.rateLimitDelayedMs);
+    }
 }

driver/src/main/java/oracle/nosql/driver/ops/RetryStats.java

@@ -36,7 +36,7 @@ public class RetryStats {
      * Internal use only.
      * Create a new retry stats object.
      */
-    RetryStats() {
+    public RetryStats() {
         this.retries = 0;
         this.delayMs = 0;
         this.exceptionMap = new HashMap<Class<? extends Throwable>, Integer>();
@@ -117,6 +117,50 @@ public void clear() {
         exceptionMap.clear();
     }
 
+    public Map<Class<? extends Throwable>, Integer> getExceptionMap() {
+        return exceptionMap;
+    }
+
+    /**
+     * @hidden
+     * Internal use only.
+     * Adds stats to the current object.
+     */
+    public void addStats(RetryStats rs) {
+        if (rs == null) {
+            return;
+        }
+        delayMs += rs.getDelayMs();
+        retries += rs.getRetries();
+        Map<Class<? extends Throwable>, Integer> emap = rs.getExceptionMap();
+        if (emap == null || emap.isEmpty()) {
+            return;
+        }
+        for (Map.Entry<Class<? extends Throwable>, Integer> entry:
+             emap.entrySet()) {
+            int i = entry.getValue().intValue();
+            Integer val = exceptionMap.get(entry.getKey());
+            if (val != null) {
+                i += val.intValue();
+            }
+            exceptionMap.put(entry.getKey(), i);
+        }
+    }
+
+    @Override
+    public boolean equals(Object o) {
+        if (!(o instanceof RetryStats)) {
+            return false;
+        }
+        RetryStats rs = (RetryStats)o;
+        if (retries != rs.retries ||
+            delayMs != rs.delayMs ||
+            exceptionMap.equals(rs.exceptionMap) == false) {
+            return false;
+        }
+        return true;
+    }
+
     @Override
     public String toString() {
         StringBuilder sb = new StringBuilder();

driver/src/main/java/oracle/nosql/driver/query/QueryDriver.java

@@ -217,6 +217,8 @@ private void setQueryResult(QueryResult result) {
         result.setReadKB(theRCB.getReadKB());
         result.setReadUnits(theRCB.getReadUnits());
         result.setWriteKB(theRCB.getWriteKB());
+        result.setRateLimitDelayedMs(theRCB.getRateLimitDelayedMs());
+        result.setRetryStats(theRCB.getRetryStats());
 
         theResults = null;
         theRCB.resetKBConsumption();

driver/src/main/java/oracle/nosql/driver/query/ReceiveIter.java

@@ -16,6 +16,7 @@
 import java.util.List;
 import java.util.TreeSet;
 
+import oracle.nosql.driver.NoSQLException;
 import oracle.nosql.driver.RetryableException;
 import oracle.nosql.driver.ops.QueryRequest;
 import oracle.nosql.driver.ops.QueryResult;
@@ -390,6 +391,36 @@ private boolean sortingNext(
         }
     }
 
+    /**
+     * Execute a copy of a request.
+     * After execution, copy specific fields from the request copy back
+     * to the original request. This is needed because client.execute() may
+     * set or update specific fields in the request, such as startTime,
+     * retry stats, rate limiters, etc.
+     */
+    private QueryResult execute(RuntimeControlBlock rcb,
+                                QueryRequest origRequest,
+                                QueryRequest reqCopy) {
+        NoSQLException e = null;
+        QueryResult result = null;
+        try {
+            result = (QueryResult)rcb.getClient().execute(reqCopy);
+        } catch (NoSQLException qe) {
+            e = qe;
+        }
+        /*
+         * Copy values back to original request, even when the execute()
+         * throws an error. This is because even in the error case the
+         * execute() call may have set/updated these values.
+         */
+        reqCopy.copyTo(origRequest);
+        /* if the execute() call threw an error, throw it here */
+        if (e != null) {
+            throw e;
+        }
+        return result;
+    }
+
     /*
      * Make sure we receive (and cache) at least one result per partition
      * (except from partitions that do not contain any results at all).
@@ -405,15 +436,16 @@ private void initPartitionSort(
          * the partition whose id is specified in theContinuationKey and
          * from any other partition that is co-located with that partition.
          */
-        QueryRequest req = rcb.getRequest().copyInternal();
-        req.setContKey(state.theContinuationKey);
+        QueryRequest origRequest = rcb.getRequest();
+        QueryRequest reqCopy = origRequest.copyInternal();
+        reqCopy.setContKey(state.theContinuationKey);
 
         if (rcb.getTraceLevel() >= 1) {
             rcb.trace("ReceiveIter : executing remote request for " +
                       "sorting phase 1");
         }
 
-        QueryResult result = (QueryResult)rcb.getClient().execute(req);
+        QueryResult result = execute(rcb, origRequest, reqCopy);
 
         int numPids = result.getNumPids();
         List<MapValue> results = result.getResultsInternal();
@@ -423,6 +455,8 @@ private void initPartitionSort(
         rcb.tallyReadKB(result.getReadKB());
         rcb.tallyReadUnits(result.getReadUnits());
         rcb.tallyWriteKB(result.getWriteKB());
+        rcb.tallyRateLimitDelayedMs(result.getRateLimitDelayedMs());
+        rcb.tallyRetryStats(result.getRetryStats());
 
         if (rcb.getTraceLevel() >= 1) {
             rcb.trace("ReceiveIter.initPartitionSort() : got result.\n" +
@@ -708,31 +742,32 @@ MapValue next() {
 
         void fetch() {
 
-            QueryRequest req = theRCB.getRequest().copyInternal();
-            req.setContKey(theContinuationKey);
-            req.setShardId(theIsForShard ? theShardOrPartId : -1);
+            QueryRequest origRequest = theRCB.getRequest();
+            QueryRequest reqCopy = origRequest.copyInternal();
+            reqCopy.setContKey(theContinuationKey);
+            reqCopy.setShardId(theIsForShard ? theShardOrPartId : -1);
 
             if (doesSort() && !theIsForShard) {
                 theState.theMemoryConsumption -= theResultsSize;
                 theRCB.decMemoryConsumption(theResultsSize);
                 long numResults =
-                    ((req.getMaxMemoryConsumption() - theState.theDupElimMemory) /
+                    ((reqCopy.getMaxMemoryConsumption() - theState.theDupElimMemory) /
                      ((theState.theSortedScanners.size() + 1) *
                       (theState.theTotalResultsSize /
                        theState.theTotalNumResults)));
                 if (numResults > 2048) {
                     numResults = 2048;
                 }
-                req.setLimit((int)numResults);
+                reqCopy.setLimit((int)numResults);
             }
 
             if (theRCB.getTraceLevel() >= 1) {
                 theRCB.trace("RemoteScanner : executing remote request. spid = " +
                              theShardOrPartId);
-                assert(req.hasDriver());
+                assert(reqCopy.hasDriver());
             }
 
-            QueryResult result = (QueryResult)theRCB.getClient().execute(req);
+            QueryResult result = execute(theRCB, origRequest, reqCopy);
 
             theResults = result.getResultsInternal();
             theContinuationKey = result.getContinuationKey();
@@ -742,6 +777,8 @@ void fetch() {
             theRCB.tallyReadKB(result.getReadKB());
             theRCB.tallyReadUnits(result.getReadUnits());
             theRCB.tallyWriteKB(result.getWriteKB());
+            theRCB.tallyRateLimitDelayedMs(result.getRateLimitDelayedMs());
+            theRCB.tallyRetryStats(result.getRetryStats());
 
             assert(result.reachedLimit() || !theMoreRemoteResults);
 

driver/src/main/java/oracle/nosql/driver/query/RuntimeControlBlock.java

@@ -12,6 +12,7 @@
 import oracle.nosql.driver.Consistency;
 import oracle.nosql.driver.http.Client;
 import oracle.nosql.driver.ops.QueryRequest;
+import oracle.nosql.driver.ops.RetryStats;
 import oracle.nosql.driver.values.FieldValue;
 
 /**
@@ -62,6 +63,12 @@ public class RuntimeControlBlock {
     private int theReadUnits;
     private int theWriteKB;
 
+    /*
+     * Stats for retries and rate limiting during query batch execution
+     */
+    private int theRateLimitDelayedMs;
+    private RetryStats theRetryStats;
+
     /*
      * The number of memory bytes consumed by the query at the client for
      * blocking operations (duplicate eleimination, sorting). Not applicable
@@ -179,6 +186,20 @@ public void setRegVal(int regId, FieldValue value) {
         theRegisters[regId] = value;
     }
 
+    public void tallyRateLimitDelayedMs(int ms) {
+        theRateLimitDelayedMs += ms;
+    }
+
+    public void tallyRetryStats(RetryStats rs) {
+        if (rs == null) {
+            return;
+        }
+        if (theRetryStats == null) {
+            theRetryStats = new RetryStats();
+        }
+        theRetryStats.addStats(rs);
+    }
+
     public void tallyReadKB(int nkb) {
         theReadKB += nkb;
     }
@@ -203,12 +224,22 @@ public void resetKBConsumption() {
         theReadKB = 0;
         theReadUnits = 0;
         theWriteKB = 0;
+        theRateLimitDelayedMs = 0;
+        theRetryStats = null;
     }
 
     public int getWriteKB() {
         return theWriteKB;
     }
 
+    public RetryStats getRetryStats() {
+        return theRetryStats;
+    }
+
+    public int getRateLimitDelayedMs() {
+        return theRateLimitDelayedMs;
+    }
+
     public void setReachedLimit(boolean value) {
         theReachedLimit = value;
     }