Remove the requirement of carrying batchCount and args in entry.exit() (#3114)

sentinel-core/src/main/java/com/alibaba/csp/sentinel/AsyncEntry.java

@@ -35,6 +35,10 @@ public class AsyncEntry extends CtEntry {
         super(resourceWrapper, chain, context);
     }
 
+    AsyncEntry(ResourceWrapper resourceWrapper, ProcessorSlot<Object> chain, Context context, int count, Object[] args) {
+        super(resourceWrapper, chain, context, count, args);
+    }
+
     /**
      * Remove current entry from local context, but does not exit.
      */

sentinel-core/src/main/java/com/alibaba/csp/sentinel/CtEntry.java

@@ -42,7 +42,11 @@ class CtEntry extends Entry {
     protected LinkedList<BiConsumer<Context, Entry>> exitHandlers;
 
     CtEntry(ResourceWrapper resourceWrapper, ProcessorSlot<Object> chain, Context context) {
-        super(resourceWrapper);
+        this(resourceWrapper, chain, context, 1, OBJECTS0);
+    }
+
+    CtEntry(ResourceWrapper resourceWrapper, ProcessorSlot<Object> chain, Context context, int count, Object[] args) {
+        super(resourceWrapper, count, args);
         this.chain = chain;
         this.context = context;
 

sentinel-core/src/main/java/com/alibaba/csp/sentinel/CtSph.java

@@ -86,7 +86,7 @@ public class CtSph implements Sph {
             return asyncEntryWithNoChain(resourceWrapper, context);
         }
 
-        AsyncEntry asyncEntry = new AsyncEntry(resourceWrapper, chain, context);
+        AsyncEntry asyncEntry = new AsyncEntry(resourceWrapper, chain, context, count, args);
         try {
             chain.entry(context, resourceWrapper, null, count, prioritized, args);
             // Initiate the async context only when the entry successfully passed the slot chain.
@@ -143,7 +143,7 @@ public class CtSph implements Sph {
             return new CtEntry(resourceWrapper, null, context);
         }
 
-        Entry e = new CtEntry(resourceWrapper, chain, context);
+        Entry e = new CtEntry(resourceWrapper, chain, context, count, args);
         try {
             chain.entry(context, resourceWrapper, null, count, prioritized, args);
         } catch (BlockException e1) {

sentinel-core/src/main/java/com/alibaba/csp/sentinel/Entry.java

@@ -53,7 +53,7 @@ import com.alibaba.csp.sentinel.context.Context;
  */
 public abstract class Entry implements AutoCloseable {
 
-    private static final Object[] OBJECTS0 = new Object[0];
+    protected static final Object[] OBJECTS0 = new Object[0];
 
     private final long createTimestamp;
     private long completeTimestamp;
@@ -69,9 +69,19 @@ public abstract class Entry implements AutoCloseable {
 
     protected final ResourceWrapper resourceWrapper;
 
+    protected final int count;
+
+    protected final Object[] args;
+
     public Entry(ResourceWrapper resourceWrapper) {
+        this(resourceWrapper, 1, OBJECTS0);
+    }
+
+    public Entry(ResourceWrapper resourceWrapper, int count, Object[] args) {
         this.resourceWrapper = resourceWrapper;
         this.createTimestamp = TimeUtil.currentTimeMillis();
+        this.count = count;
+        this.args = args;
     }
 
     public ResourceWrapper getResourceWrapper() {
@@ -80,15 +90,15 @@ public abstract class Entry implements AutoCloseable {
 
     /**
      * Complete the current resource entry and restore the entry stack in context.
-     *
+     * Do not need to carry count or args parameter, initialization does
      * @throws ErrorEntryFreeException if entry in current context does not match current entry
      */
     public void exit() throws ErrorEntryFreeException {
-        exit(1, OBJECTS0);
+        exit(count, args);
     }
 
     public void exit(int count) throws ErrorEntryFreeException {
-        exit(count, OBJECTS0);
+        exit(count, args);
     }
 
     /**

sentinel-core/src/main/java/com/alibaba/csp/sentinel/SphO.java

@@ -221,6 +221,6 @@ public class SphO {
     }
 
     public static void exit() {
-        exit(1, OBJECTS0);
+        ContextUtil.getContext().getCurEntry().exit();
     }
 }

sentinel-core/src/test/java/com/alibaba/csp/sentinel/SphUTest.java

@@ -157,4 +157,14 @@ public class SphUTest {
 
         e.exit(2, arg0, arg1);
     }
+
+    @Test
+    public void testEntryExitAutomation() throws BlockException{
+        String[] args = {"foo", "baz"};
+        int batchCount = 3;
+        Entry e = SphU.entry("testEntryExitAutomation", EntryType.IN, 3, args);
+        e.exit();
+        // The number of success is automatically updated based on batchCount when exit
+        assertEquals(batchCount, e.getCurNode().totalSuccess());
+    }
 }