Merge pull request #2216 from AnuGayan/master-temp2

Update Throttler to support increase counter by more than one
wso2 · Sep 18, 2024 · 66aacb6 · 66aacb6
2 parents 42637cb + 59d3f19
commit 66aacb6
Show file tree

Hide file tree

Showing 2 changed files with 29 additions and 7 deletions.
diff --git a/modules/commons/src/main/java/org/apache/synapse/commons/throttle/core/CallerContext.java b/modules/commons/src/main/java/org/apache/synapse/commons/throttle/core/CallerContext.java
@@ -138,10 +138,11 @@ private void initAccess(CallerConfiguration configuration, ThrottleContext throt
      * @param configuration   -  The Configuration for this caller
      * @param throttleContext -The Throttle Context
      * @param currentTime     -The system current time
+     * @param eventCount      -The event count
      * @return boolean        -The boolean value which say access will allow or not
      */
     private boolean canAccessIfUnitTimeNotOver(CallerConfiguration configuration,
-                                               ThrottleContext throttleContext, long currentTime) {
+                                               ThrottleContext throttleContext, long currentTime, Long eventCount) {
         boolean canAccess = false;
         int maxRequest = configuration.getMaximumRequestPerUnitTime();
         if (maxRequest != 0) {
@@ -154,7 +155,7 @@ private boolean canAccessIfUnitTimeNotOver(CallerConfiguration configuration,
                             + configuration.getID() + " nextAccessTime=" + this.nextAccessTime);
                 }
                 canAccess = true;     // can continue access
-                this.localCount.incrementAndGet();
+                this.localCount.addAndGet(eventCount);
                 // Send the current state to others (clustered env)
                 throttleContext.flushCallerContext(this, id);
                 // can complete access
@@ -209,7 +210,7 @@ private boolean canAccessIfUnitTimeNotOver(CallerConfiguration configuration,
 
                         if (log.isDebugEnabled()) {
                             log.debug("Caller=" + this.getId() + " has reset counters and added for replication when unit "
-                                      + "time is not over");
+                                    + "time is not over");
                         }
                     } else {
                         if (log.isDebugEnabled()) {
@@ -373,6 +374,21 @@ public void cleanUpCallers(CallerConfiguration configuration,
      */
     public boolean canAccess(ThrottleContext throttleContext, CallerConfiguration configuration,
                              long currentTime) throws ThrottleException {
+        return canAccess(throttleContext, configuration, currentTime, 1L);
+    }
+
+    /**
+     * Check whether that caller can access or not, based on current state and pre-defined policy
+     *
+     * @param throttleContext -The Context for this caller - runtime state
+     * @param configuration   -The Configuration for this caller - data from policy
+     * @param currentTime     -The current system time
+     * @param eventCount      -The event count
+     * @return boolean        -The boolean value which say access will allow or not
+     * @throws ThrottleException throws for invalid throttle configuration
+     */
+    public boolean canAccess(ThrottleContext throttleContext, CallerConfiguration configuration,
+                             long currentTime, Long eventCount) throws ThrottleException {
         RequestContext requestContext = new RequestContext(currentTime);
         boolean canAccess;
         if (configuration == null) {
@@ -407,14 +423,15 @@ public boolean canAccess(ThrottleContext throttleContext, CallerConfiguration co
                 log.debug("LATENCY FOR THROTTLE PROCESSING: " + duration + " ms");
             }
         } else {
-            canAccess = canAccessBasedOnUnitTime(configuration, throttleContext, currentTime);
+            canAccess = canAccessBasedOnUnitTime(configuration, throttleContext, currentTime, eventCount);
         }
         return canAccess;
     }
 
-    private boolean canAccessBasedOnUnitTime(CallerConfiguration configuration, ThrottleContext throttleContext, long currentTime) {
+    private boolean canAccessBasedOnUnitTime(CallerConfiguration configuration, ThrottleContext throttleContext,
+                                             long currentTime, Long eventCount) {
         if (this.nextTimeWindow > currentTime) {
-            return canAccessIfUnitTimeNotOver(configuration, throttleContext, currentTime);
+            return canAccessIfUnitTimeNotOver(configuration, throttleContext, currentTime, eventCount);
         } else {
             return canAccessIfUnitTimeOver(configuration, throttleContext, currentTime);
         }

diff --git a/...src/main/java/org/apache/synapse/commons/throttle/core/RoleBasedAccessRateController.java b/...src/main/java/org/apache/synapse/commons/throttle/core/RoleBasedAccessRateController.java
@@ -57,6 +57,11 @@ public RoleBasedAccessRateController() {
      */
     public AccessInformation canAccess(ThrottleContext throttleContext,
                                        String consumerKey, String roleID) throws ThrottleException {
+        return canAccess(throttleContext, consumerKey, roleID, 1L);
+    }
+
+    public AccessInformation canAccess(ThrottleContext throttleContext,
+                                       String consumerKey, String roleID, Long eventCount) throws ThrottleException {
 
         String type = "role";
 
@@ -116,7 +121,7 @@ public AccessInformation canAccess(ThrottleContext throttleContext,
                 if (caller != null) {
                     long currentTime = System.currentTimeMillis();
 
-                    if (!caller.canAccess(throttleContext, configuration, currentTime)) {
+                    if (!caller.canAccess(throttleContext, configuration, currentTime, eventCount)) {
                         //if current caller cannot access , then perform cleaning
                         log.info(ACCESS_DENIED_TEMPORALLY);
                         throttleContext.processCleanList(currentTime);