Minor refactoring, reduce dup code

0xaa4eb · Jul 8, 2024 · 74efc44 · 74efc44
1 parent 358292a
commit 74efc44
Show file tree

Hide file tree

Showing 5 changed files with 31 additions and 41 deletions.
diff --git a/ulyp-agent-core/src/main/java/com/ulyp/agent/Recorder.java b/ulyp-agent-core/src/main/java/com/ulyp/agent/Recorder.java
@@ -17,6 +17,7 @@
 import java.util.stream.Collectors;
 import java.util.stream.Stream;
 
+import org.jetbrains.annotations.NotNull;
 import org.jetbrains.annotations.Nullable;
 import org.jetbrains.annotations.TestOnly;
 
@@ -26,7 +27,6 @@
 public class Recorder {
 
     public static final AtomicInteger recordingIdGenerator = new AtomicInteger(-1);
-
     /**
     * Keeps current active recording session count. Based on the fact that most of the time there is no
     * recording sessions and this counter is equal to 0, it's possible to make a small performance optimization.
@@ -86,57 +86,47 @@ public void enableRecording() {
         }
     }
 
-    public int startOrContinueRecordingOnMethodEnter(int methodId, @Nullable Object callee, Object[] args) {
+    public int startRecordingOnMethodEnter(int methodId, @Nullable Object callee, Object[] args) {
         if (startRecordingPolicy.canStartRecording()) {
-            RecordingState recordingState = threadLocalRecordingState.get();
-            if (recordingState == null) {
-                recordingState = new RecordingState();
-                RecordingMetadata recordingMetadata = generateRecordingMetadata();
-                recordingState.setRecordingMetadata(recordingMetadata);
-                recordingState.setEnabled(false);
-
-                threadLocalRecordingState.set(recordingState);
-
-                currentRecordingSessionCount.incrementAndGet();
-                if (LoggingSettings.DEBUG_ENABLED) {
-                    log.debug("Started recording {} at method {}", recordingMetadata.getId(), methodRepository.get(methodId));
-                }
-                recordingsCounter.inc();
-                recordingState.setEnabled(true);
-                recordingEventQueue.enqueueRecordingStarted(recordingMetadata);
-            }
+            RecordingState recordingState = initializeRecordingState(methodId);
 
             return onMethodEnter(recordingState, methodId, callee, args);
         } else {
             return -1;
         }
     }
 
-    public int startOrContinueRecordingOnConstructorEnter(int methodId, Object[] args) {
+    public int startRecordingOnConstructorEnter(int methodId, Object[] args) {
         if (startRecordingPolicy.canStartRecording()) {
-            RecordingState recordingState = threadLocalRecordingState.get();
-            if (recordingState == null) {
-                recordingState = new RecordingState();
-                recordingState.setEnabled(false);
-                RecordingMetadata recordingMetadata = generateRecordingMetadata();
-                recordingState.setRecordingMetadata(recordingMetadata);
-                threadLocalRecordingState.set(recordingState);
-
-                currentRecordingSessionCount.incrementAndGet();
-                if (LoggingSettings.DEBUG_ENABLED) {
-                    log.debug("Started recording {} at method {}", recordingMetadata.getId(), methodRepository.get(methodId).toShortString());
-                }
-                recordingsCounter.inc();
-                recordingState.setEnabled(true);
-                recordingEventQueue.enqueueRecordingStarted(recordingMetadata);
-            }
+            RecordingState recordingState = initializeRecordingState(methodId);
 
             return onConstructorEnter(recordingState, methodId, args);
         } else {
             return -1;
         }
     }
 
+    @NotNull
+    private RecordingState initializeRecordingState(int methodId) {
+        RecordingState recordingState = threadLocalRecordingState.get();
+        if (recordingState == null) {
+            recordingState = new RecordingState();
+            recordingState.setEnabled(false);
+            RecordingMetadata recordingMetadata = generateRecordingMetadata();
+            recordingState.setRecordingMetadata(recordingMetadata);
+            threadLocalRecordingState.set(recordingState);
+
+            currentRecordingSessionCount.incrementAndGet();
+            if (LoggingSettings.DEBUG_ENABLED) {
+                log.debug("Started recording {} at method {}", recordingMetadata.getId(), methodRepository.get(methodId));
+            }
+            recordingsCounter.inc();
+            recordingState.setEnabled(true);
+            recordingEventQueue.enqueueRecordingStarted(recordingMetadata);
+        }
+        return recordingState;
+    }
+
     public int onConstructorEnter(int methodId, Object[] args) {
         return onMethodEnter(threadLocalRecordingState.get(), methodId, null, args);
     }

diff --git a/ulyp-agent-core/src/test/java/com/ulyp/agent/RecorderCurrentSessionCountTest.java b/ulyp-agent-core/src/test/java/com/ulyp/agent/RecorderCurrentSessionCountTest.java
@@ -81,7 +81,7 @@ public void run() {
             X callee = new X();
 
             for (int i = 0; i < recordingsCount && !Thread.currentThread().isInterrupted(); i++) {
-                int callId = recorder.startOrContinueRecordingOnMethodEnter(methodIdx, callee, new Object[5]);
+                int callId = recorder.startRecordingOnMethodEnter(methodIdx, callee, new Object[5]);
 
                 Assertions.assertTrue(Recorder.currentRecordingSessionCount.get() > 0, "Since at least one recording session is active, " +
                     "Recorder.currentRecordingSessionCount must be positive");

diff --git a/ulyp-agent-core/src/test/java/com/ulyp/agent/RecorderTest.java b/ulyp-agent-core/src/test/java/com/ulyp/agent/RecorderTest.java
@@ -65,7 +65,7 @@ public void tearDown() {
     @Test
     void shouldRecordDataWhenRecordingIsFinished() throws InterruptedException, TimeoutException {
         X recorded = new X();
-        int callId = recorder.startOrContinueRecordingOnMethodEnter(methodIdx, recorded, new Object[] {5});
+        int callId = recorder.startRecordingOnMethodEnter(methodIdx, recorded, new Object[] {5});
         recorder.onMethodExit(methodIdx, "ABC", null, callId);
         callRecordQueue.sync(Duration.ofSeconds(5));
 
@@ -78,7 +78,7 @@ void testTemporaryRecordingDisableWithOngoingRecording() throws InterruptedExcep
         Recorder.recordingIdGenerator.set(0);
 
         X recorded = new X();
-        int callId1 = recorder.startOrContinueRecordingOnMethodEnter(methodIdx, recorded, new Object[] {5});
+        int callId1 = recorder.startRecordingOnMethodEnter(methodIdx, recorded, new Object[] {5});
 
         recorder.disableRecording();
 

diff --git a/ulyp-agent/src/main/java/com/ulyp/agent/ConstructorCallRecordingAdvice.java b/ulyp-agent/src/main/java/com/ulyp/agent/ConstructorCallRecordingAdvice.java
@@ -23,7 +23,7 @@ static void enter(
 
         // This if check is ugly, but the code is wired into bytecode, so it's more efficient to check right away instead of calling a method
         if (methodId >= MethodRepository.RECORD_METHODS_MIN_ID) {
-            callId = RecorderInstance.instance.startOrContinueRecordingOnConstructorEnter(methodId, arguments);
+            callId = RecorderInstance.instance.startRecordingOnConstructorEnter(methodId, arguments);
         } else {
             if (Recorder.currentRecordingSessionCount.get() > 0 && RecorderInstance.instance.recordingIsActiveInCurrentThread()) {
                 callId = RecorderInstance.instance.onConstructorEnter(methodId, arguments);

diff --git a/ulyp-agent/src/main/java/com/ulyp/agent/MethodCallRecordingAdvice.java b/ulyp-agent/src/main/java/com/ulyp/agent/MethodCallRecordingAdvice.java
@@ -26,7 +26,7 @@ static void enter(
         if (methodId >= MethodRepository.RECORD_METHODS_MIN_ID) {
 
             // noinspection UnusedAssignment local variable callId is used by exit() method
-            callId = RecorderInstance.instance.startOrContinueRecordingOnMethodEnter(methodId, callee, arguments);
+            callId = RecorderInstance.instance.startRecordingOnMethodEnter(methodId, callee, arguments);
         } else {
 
             if (Recorder.currentRecordingSessionCount.get() > 0 && RecorderInstance.instance.recordingIsActiveInCurrentThread()) {