Merge pull request #12658 from nisan-abeywickrama/retry-after-header

Add nextAccess timestamp to synapse context when throttle by burst control
wso2 · Nov 12, 2024 · 0f49253 · 0f49253
2 parents 907fb00 + 0f11d8b
commit 0f49253
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 31 deletions.
diff --git a/...way/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java b/...way/src/main/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandler.java
@@ -42,6 +42,7 @@
 import org.apache.synapse.commons.throttle.core.ThrottleContext;
 import org.apache.synapse.commons.throttle.core.ThrottleException;
 import org.apache.synapse.commons.throttle.core.ThrottleFactory;
+import org.apache.synapse.commons.throttle.core.CallerContext;
 import org.apache.synapse.commons.throttle.core.factory.ThrottleContextFactory;
 import org.apache.synapse.core.SynapseEnvironment;
 import org.apache.synapse.core.axis2.Axis2MessageContext;
@@ -1237,6 +1238,9 @@ public boolean isSubscriptionLevelSpike(MessageContext synCtx, String throttleKe
                 }
 
                 if (info != null && !info.isAccessAllowed()) {
+                    CallerContext callerContext = subscriptionLevelSpikeArrestThrottleContext.getCallerContext(throttleKey);
+                    long timestamp = callerContext.getNextAccessTime();
+                    synCtx.setProperty(APIThrottleConstants.THROTTLED_NEXT_ACCESS_TIMESTAMP, timestamp);
                     synCtx.setProperty(APIThrottleConstants.THROTTLED_OUT_REASON, APIThrottleConstants.SUBSCRIPTON_BURST_LIMIT_EXCEEDED);
                     log.debug("Subscription level burst control limit exceeded for key " + throttleKey);
                     return true;

diff --git a/...src/test/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandlerTest.java b/...src/test/java/org/wso2/carbon/apimgt/gateway/handlers/throttling/ThrottleHandlerTest.java
@@ -472,37 +472,6 @@ public void testMsgThrottleOutWhenCustomThrottlingLimitExceeded() {
         throttleDataHolder.removeKeyTemplate("$user");
     }
 
-    @Test
-    public void testMsgThrottleOutWhenHittingSubscriptionLevelSpike() {
-        ThrottleDataHolder throttleDataHolder = ThrottleDataHolder.getInstance();
-
-        ThrottleHandler throttleHandler = new ThrottlingHandlerWrapper(timer, throttleDataHolder, throttleEvaluator,
-                accessInformation);
-        throttleHandler.setSandboxMaxCount("100");
-        SynapseEnvironment synapseEnvironment = Mockito.mock(SynapseEnvironment.class);
-        throttleHandler.init(synapseEnvironment);
-
-        MessageContext messageContext = TestUtils.getMessageContextWithAuthContext(apiContext, apiVersion);
-        messageContext.setProperty(VERB_INFO_DTO, verbInfoDTO);
-        ((Axis2MessageContext) messageContext).getAxis2MessageContext().getProperty(org.apache.axis2.context
-                .MessageContext.TRANSPORT_HEADERS);
-        AuthenticationContext authenticationContext = (AuthenticationContext) messageContext.getProperty
-                (API_AUTH_CONTEXT);
-        authenticationContext.setApiTier(throttlingTier);
-        authenticationContext.setKeyType("SANDBOX");
-        authenticationContext.setSpikeArrestLimit(100);
-        authenticationContext.setStopOnQuotaReach(true);
-
-        messageContext.setProperty(API_AUTH_CONTEXT, authenticationContext);
-
-        verbInfo.setConditionGroups(conditionGroupDTOs);
-        ArrayList<ConditionGroupDTO> matchingConditions = new ArrayList<>();
-        matchingConditions.add(conditionGroupDTO);
-        throttleDataHolder.addKeyTemplate("$user", "$user");
-        Mockito.when(accessInformation.isAccessAllowed()).thenReturn(false);
-        Assert.assertFalse(throttleHandler.handleRequest(messageContext));
-    }
-
     @Test
     public void testHandleResponse() {
         ThrottleDataHolder throttleDataHolder = ThrottleDataHolder.getInstance();