From 4c3b592fb09d3b8acaa2fe4222d3f6a2fdeb8cd7 Mon Sep 17 00:00:00 2001
From: Gaston Thea <gaston.thea@harness.io>
Date: Tue, 10 Dec 2024 17:34:40 -0300
Subject: [PATCH] Update recorder sync helper to handle multiple listeners

---
 .../synchronizer/RecorderSyncHelperImpl.java  | 23 +++++++++++-----
 .../RecorderSyncHelperImplTest.java           | 27 +++++++++++++++++++
 2 files changed, 44 insertions(+), 6 deletions(-)

diff --git a/src/main/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImpl.java b/src/main/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImpl.java
index a2582584a..4351930fb 100644
--- a/src/main/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImpl.java
+++ b/src/main/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImpl.java
@@ -5,6 +5,8 @@
 import androidx.annotation.NonNull;
 
 import java.lang.ref.WeakReference;
+import java.util.HashSet;
+import java.util.Set;
 import java.util.concurrent.atomic.AtomicInteger;
 import java.util.concurrent.atomic.AtomicLong;
 
@@ -26,7 +28,7 @@ public class RecorderSyncHelperImpl<T extends InBytesSizable> implements Recorde
     private final int mMaxQueueSize;
     private final long mMaxQueueSizeInBytes;
     private final SplitTaskType mTaskType;
-    private WeakReference<SplitTaskExecutionListener> mTaskExecutionListener;
+    private final Set<WeakReference<SplitTaskExecutionListener>> mTaskExecutionListener;
 
     public RecorderSyncHelperImpl(SplitTaskType taskType,
                                   StoragePusher<T> storage,
@@ -40,7 +42,7 @@ public RecorderSyncHelperImpl(SplitTaskType taskType,
         mTotalPushedSizeInBytes = new AtomicLong(0);
         mMaxQueueSize = maxQueueSize;
         mMaxQueueSizeInBytes = maxQueueSizeInBytes;
-        mTaskExecutionListener = new WeakReference<>(null);
+        mTaskExecutionListener = new HashSet<>();
     }
 
     @Override
@@ -67,19 +69,28 @@ public void taskExecuted(@NonNull SplitTaskExecutionInfo taskInfo) {
                     SplitTaskExecutionInfo.NON_SENT_BYTES));
         }
 
-        if (mTaskExecutionListener.get() != null) {
-            mTaskExecutionListener.get().taskExecuted(taskInfo);
+        for (WeakReference<SplitTaskExecutionListener> reference : mTaskExecutionListener) {
+            SplitTaskExecutionListener listener = reference.get();
+            if (listener != null) {
+                listener.taskExecuted(taskInfo);
+            }
         }
     }
 
     @Override
     public void addListener(SplitTaskExecutionListener listener) {
-        mTaskExecutionListener = new WeakReference<>(listener);
+        mTaskExecutionListener.add(new WeakReference<>(listener));
     }
 
     @Override
     public void removeListener(SplitTaskExecutionListener listener) {
-        mTaskExecutionListener = new WeakReference<>(null);
+        for (WeakReference<SplitTaskExecutionListener> reference : mTaskExecutionListener) {
+            SplitTaskExecutionListener listenerRef = reference.get();
+            if (listenerRef != null && listenerRef.equals(listener)) {
+                mTaskExecutionListener.remove(reference);
+                break;
+            }
+        }
     }
 
     private void pushAsync(T entity) {
diff --git a/src/test/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImplTest.java b/src/test/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImplTest.java
index dcd5a45ad..1888bf831 100644
--- a/src/test/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImplTest.java
+++ b/src/test/java/io/split/android/client/service/synchronizer/RecorderSyncHelperImplTest.java
@@ -52,13 +52,40 @@ public void listenerIsNotCalledOnceRemoved() {
                 mTaskExecutor);
 
         SplitTaskExecutionListener listener = mock(SplitTaskExecutionListener.class);
+        SplitTaskExecutionListener listener2 = mock(SplitTaskExecutionListener.class);
         SplitTaskExecutionInfo executionInfo = SplitTaskExecutionInfo.success(SplitTaskType.IMPRESSIONS_RECORDER);
 
         syncHelper.addListener(listener);
+        syncHelper.addListener(listener2);
         syncHelper.taskExecuted(executionInfo);
         syncHelper.removeListener(listener);
         syncHelper.taskExecuted(executionInfo);
 
         verify(listener, times(1)).taskExecuted(executionInfo);
+        verify(listener2, times(2)).taskExecuted(executionInfo);
+    }
+
+    @Test
+    public void multipleListenersCanBeAdded() {
+        mStoragePusher = mock(StoragePusher.class);
+        mTaskExecutor = mock(SplitTaskExecutor.class);
+
+        RecorderSyncHelperImpl<KeyImpression> syncHelper = new RecorderSyncHelperImpl<>(
+                SplitTaskType.IMPRESSIONS_RECORDER,
+                mStoragePusher,
+                1,
+                1,
+                mTaskExecutor);
+
+        SplitTaskExecutionListener listener1 = mock(SplitTaskExecutionListener.class);
+        SplitTaskExecutionListener listener2 = mock(SplitTaskExecutionListener.class);
+        SplitTaskExecutionInfo executionInfo = SplitTaskExecutionInfo.success(SplitTaskType.IMPRESSIONS_RECORDER);
+
+        syncHelper.addListener(listener1);
+        syncHelper.addListener(listener2);
+        syncHelper.taskExecuted(executionInfo);
+
+        verify(listener1).taskExecuted(executionInfo);
+        verify(listener2).taskExecuted(executionInfo);
     }
 }