From 153aaecf0fb4fd1fcdbb9b2615c24c9a1311647f Mon Sep 17 00:00:00 2001
From: Santiago Mola <santiago.mola@datadoghq.com>
Date: Tue, 25 Feb 2025 14:12:11 +0100
Subject: [PATCH] Add more tests to NonBlockingSemaphore and make boolean
 semaphore release atomic

---
 .../iast/util/NonBlockingSemaphore.java       |  2 +-
 .../iast/util/NonBlockingSemaphoreTest.groovy | 80 +++++++++++++++----
 2 files changed, 67 insertions(+), 15 deletions(-)

diff --git a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java
index ed96b1192dc..ef6e529980a 100644
--- a/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java
+++ b/dd-java-agent/agent-iast/src/main/java/com/datadog/iast/util/NonBlockingSemaphore.java
@@ -66,7 +66,7 @@ public boolean acquire(final int count) {
     @Override
     public int release(final int count) {
       reset();
-      return available();
+      return 1;
     }
 
     @Override
diff --git a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy
index 09bf6ad2abe..a8247cac497 100644
--- a/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy
+++ b/dd-java-agent/agent-iast/src/test/groovy/com/datadog/iast/util/NonBlockingSemaphoreTest.groovy
@@ -1,7 +1,8 @@
 package com.datadog.iast.util
 
-import datadog.trace.test.util.DDSpecification
+
 import groovy.transform.CompileDynamic
+import spock.lang.Specification
 
 import java.util.concurrent.Callable
 import java.util.concurrent.CountDownLatch
@@ -9,7 +10,7 @@ import java.util.concurrent.Executors
 import java.util.concurrent.TimeUnit
 
 @CompileDynamic
-class NonBlockingSemaphoreTest extends DDSpecification {
+class NonBlockingSemaphoreTest extends Specification {
 
   void 'test that the semaphore controls access to a shared resource (#permitCount)'(final int permitCount) {
     given:
@@ -71,6 +72,40 @@ class NonBlockingSemaphoreTest extends DDSpecification {
     2           | _
   }
 
+  void 'can never acquire more permits than the total'(final int permitCount) {
+    given:
+    final semaphore = NonBlockingSemaphore.withPermitCount(permitCount)
+
+    when:
+    final acquired = semaphore.acquire(permitCount+1)
+
+    then:
+    !acquired
+
+    where:
+    permitCount | _
+    1           | _
+    2           | _
+  }
+
+  void 'can perform extra releases'(final int permitCount) {
+    given:
+    final semaphore = NonBlockingSemaphore.withPermitCount(permitCount)
+
+    when:
+    for (int i = 0; i < permitCount * 2; i++) {
+      assert semaphore.release() == permitCount
+    }
+
+    then:
+    semaphore.available() == permitCount
+
+    where:
+    permitCount | _
+    1           | _
+    2           | _
+  }
+
   void 'reset helps recover when there is starvation (#permitCount)'(final int permitCount) {
     given:
     final semaphore = NonBlockingSemaphore.withPermitCount(permitCount)
@@ -97,24 +132,41 @@ class NonBlockingSemaphoreTest extends DDSpecification {
     given:
     final int threads = 100
     final semaphore = NonBlockingSemaphore.unlimited()
-    final latch = new CountDownLatch(threads)
-    final executors = Executors.newFixedThreadPool(threads)
 
     when:
     final acquired = (1..threads).collect {
-      executors.submit({
-        latch.countDown()
-        if (semaphore.acquire()) {
-          TimeUnit.MILLISECONDS.sleep(100)
-          semaphore.release()
-          return 1
-        }
-        return 0
-      } as Callable<Integer>)
-    }.collect { it.get() }.sum()
+      semaphore.acquire()? 1 : 0
+    }.collect { it }.sum()
 
     then:
     acquired == threads
     semaphore.available() == Integer.MAX_VALUE
+
+    when:
+    int availableAfterRelease = semaphore.release()
+
+    then:
+    availableAfterRelease == Integer.MAX_VALUE
+    semaphore.available() == Integer.MAX_VALUE
+
+    when:
+    semaphore.reset()
+
+    then:
+    semaphore.available() == Integer.MAX_VALUE
+  }
+
+  void 'cannot create a semaphore without at least 1 permit'() {
+    when:
+    NonBlockingSemaphore.withPermitCount(0)
+
+    then:
+    thrown(AssertionError)
+
+    when:
+    NonBlockingSemaphore.withPermitCount(-1)
+
+    then:
+    thrown(AssertionError)
   }
 }