Merge pull request cdapio#15390 from cdapio/bugfix/CDAP-20871-system-…

…worker-capacity-leak

[CDAP-20871] Reset concurrent request count in system worker when returning 429 response

cdap-app-fabric/src/main/java/io/cdap/cdap/internal/app/worker/system/SystemWorkerHttpHandlerInternal.java

@@ -100,6 +100,7 @@ public void run(FullHttpRequest request, HttpResponder responder) {
 
     if (requestProcessedCount.incrementAndGet() > requestLimit) {
       responder.sendStatus(HttpResponseStatus.TOO_MANY_REQUESTS);
+      requestProcessedCount.decrementAndGet();
       return;
     }
 

cdap-app-fabric/src/test/java/io/cdap/cdap/internal/app/worker/system/SystemWorkerServiceTest.java

@@ -198,14 +198,58 @@ public void testValidConcurrentRequests() throws Exception {
           == HttpResponseStatus.OK.code()) {
         okResponse++;
       } else if (responses.get(i).get().getResponseCode()
-          == HttpResponseStatus.TOO_MANY_REQUESTS.code()) {
+                 == HttpResponseStatus.TOO_MANY_REQUESTS.code()) {
         conflictResponse++;
       }
     }
     Assert.assertEquals(2, okResponse);
     Assert.assertEquals(concurrentRequests, okResponse + conflictResponse);
   }
 
+  @Test
+  public void testRepeatedConcurrentRequests() throws Exception {
+    InetSocketAddress addr = systemWorkerService.getBindAddress();
+    URI uri = URI.create(
+        String.format("http://%s:%s", addr.getHostName(), addr.getPort()));
+
+    RunnableTaskRequest request = RunnableTaskRequest.getBuilder(
+            SystemWorkerServiceTest.TestRunnableClass.class.getName())
+        .withParam("500").build();
+
+    String reqBody = GSON.toJson(request);
+    List<Callable<HttpResponse>> calls = new ArrayList<>();
+    int concurrentRequests = 6;
+
+    for (int i = 0; i < concurrentRequests; i++) {
+      calls.add(() -> HttpRequests.execute(
+          HttpRequest.post(uri.resolve("/v3Internal/system/run").toURL())
+              .withBody(reqBody).build(), new DefaultHttpRequestConfig(false)));
+    }
+
+    Executors.newFixedThreadPool(concurrentRequests).invokeAll(calls);
+
+    // Wait for requests to complete and retry them.
+    Thread.sleep(1000);
+
+    int okResponse = 0;
+    int conflictResponse = 0;
+    List<Future<HttpResponse>> responses = Executors.newFixedThreadPool(
+        concurrentRequests).invokeAll(calls);
+    for (int i = 0; i < concurrentRequests; i++) {
+      if (responses.get(i).get().getResponseCode()
+          == HttpResponseStatus.OK.code()) {
+        okResponse++;
+      } else if (responses.get(i).get().getResponseCode()
+                 == HttpResponseStatus.TOO_MANY_REQUESTS.code()) {
+        conflictResponse++;
+      }
+    }
+
+    Assert.assertEquals(5, okResponse);
+    Assert.assertEquals(concurrentRequests, okResponse + conflictResponse);
+  }
+
+
   @Test
   public void testInvalidConcurrentRequests() throws Exception {
     InetSocketAddress addr = systemWorkerService.getBindAddress();