Remove the requirement of carrying args in entry.exit() alibaba#3109

LiYangSir · May 4, 2023 · 8ecb6f2 · 8ecb6f2
1 parent 0af60f2
commit 8ecb6f2
Show file tree

Hide file tree

Showing 6 changed files with 35 additions and 8 deletions.
diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/AsyncEntry.java b/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 batchCount, Object[] args) {
+        super(resourceWrapper, chain, context, batchCount, args);
+    }
+
     /**
      * Remove current entry from local context, but does not exit.
      */

diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/CtEntry.java b/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 batchCount, Object[] args) {
+        super(resourceWrapper, batchCount, args);
         this.chain = chain;
         this.context = context;
 

diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/CtSph.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/CtSph.java
@@ -86,7 +86,7 @@ private AsyncEntry asyncEntryWithPriorityInternal(ResourceWrapper resourceWrappe
             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 @@ private Entry entryWithPriority(ResourceWrapper resourceWrapper, int count, bool
             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) {

diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/Entry.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/Entry.java
@@ -53,7 +53,7 @@
  */
 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 batchCount;
+
+    protected final Object[] args;
+
     public Entry(ResourceWrapper resourceWrapper) {
+        this(resourceWrapper, 1, OBJECTS0);
+    }
+
+    public Entry(ResourceWrapper resourceWrapper, int batchCount, Object[] args) {
         this.resourceWrapper = resourceWrapper;
         this.createTimestamp = TimeUtil.currentTimeMillis();
+        this.batchCount = batchCount;
+        this.args = args;
     }
 
     public ResourceWrapper getResourceWrapper() {
@@ -80,15 +90,15 @@ public ResourceWrapper getResourceWrapper() {
 
     /**
      * Complete the current resource entry and restore the entry stack in context.
-     *
+     * Do not need to carry the count and 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(batchCount, args);
     }
 
     public void exit(int count) throws ErrorEntryFreeException {
-        exit(count, OBJECTS0);
+        exit(count, args);
     }
 
     /**

diff --git a/sentinel-core/src/main/java/com/alibaba/csp/sentinel/SphO.java b/sentinel-core/src/main/java/com/alibaba/csp/sentinel/SphO.java
@@ -221,6 +221,6 @@ public static void exit(int count) {
     }
 
     public static void exit() {
-        exit(1, OBJECTS0);
+        ContextUtil.getContext().getCurEntry().exit();
     }
 }
diff --git a/sentinel-core/src/test/java/com/alibaba/csp/sentinel/SphUTest.java b/sentinel-core/src/test/java/com/alibaba/csp/sentinel/SphUTest.java
@@ -157,4 +157,13 @@ public void testMethodEntryAll() throws BlockException, NoSuchMethodException, S
 
         e.exit(2, arg0, arg1);
     }
+
+    @Test
+    public void testEntryExitAutomation() throws BlockException{
+        String[] args = {"foo", "baz"};
+        int batchCount = 3;
+        Entry e = SphU.entry("resourceName", EntryType.IN, 3, args);
+        e.exit();
+        assertEquals(batchCount, e.getCurNode().totalSuccess());
+    }
 }