2 files changed, 43 insertions, 4 deletions

diff --git a/libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java b/libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java
index 7624b693ac43..fe60037483c4 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/util/AcceptOnceConsumer.java
@@ -27,9 +27,10 @@ import java.util.function.Consumer;
  */
 public class AcceptOnceConsumer<T> implements Consumer<T> {
     private final Consumer<T> mCallback;
-    private final DataProducer<T> mProducer;
+    private final AcceptOnceProducerCallback<T> mProducer;
 
-    public AcceptOnceConsumer(@NonNull DataProducer<T> producer, @NonNull Consumer<T> callback) {
+    public AcceptOnceConsumer(@NonNull AcceptOnceProducerCallback<T> producer,
+            @NonNull Consumer<T> callback) {
         mProducer = producer;
         mCallback = callback;
     }
@@ -37,6 +38,20 @@ public class AcceptOnceConsumer<T> implements Consumer<T> {
     @Override
     public void accept(@NonNull T t) {
         mCallback.accept(t);
-        mProducer.removeDataChangedCallback(this);
+        mProducer.onConsumerReadyToBeRemoved(this);
+    }
+
+    /**
+     * Interface to allow the {@link AcceptOnceConsumer} to notify the client that created it,
+     * when it is ready to be removed. This allows the client to remove the consumer object
+     * when it deems it is safe to do so.
+     * @param <T> The type of data this callback accepts through {@link #onConsumerReadyToBeRemoved}
+     */
+    public interface AcceptOnceProducerCallback<T> {
+
+        /**
+         * Notifies that the given {@code callback} is ready to be removed
+         */
+        void onConsumerReadyToBeRemoved(Consumer<T> callback);
     }
 }
diff --git a/libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java b/libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java
index cbaa27712015..46c925aaf8a2 100644
--- a/libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java
+++ b/libs/WindowManager/Jetpack/src/androidx/window/util/BaseDataProducer.java
@@ -19,6 +19,7 @@ package androidx.window.util;
 import androidx.annotation.GuardedBy;
 import androidx.annotation.NonNull;
 
+import java.util.HashSet;
 import java.util.LinkedHashSet;
 import java.util.Optional;
 import java.util.Set;
@@ -31,11 +32,14 @@ import java.util.function.Consumer;
  *
  * @param <T> The type of data this producer returns through {@link DataProducer#getData}.
  */
-public abstract class BaseDataProducer<T> implements DataProducer<T> {
+public abstract class BaseDataProducer<T> implements DataProducer<T>,
+        AcceptOnceConsumer.AcceptOnceProducerCallback<T> {
 
     private final Object mLock = new Object();
     @GuardedBy("mLock")
     private final Set<Consumer<T>> mCallbacks = new LinkedHashSet<>();
+    @GuardedBy("mLock")
+    private final Set<Consumer<T>> mCallbacksToRemove = new HashSet<>();
 
     /**
      * Adds a callback to the set of callbacks listening for data. Data is delivered through
@@ -85,6 +89,26 @@ public abstract class BaseDataProducer<T> implements DataProducer<T> {
             for (Consumer<T> callback : mCallbacks) {
                 callback.accept(value);
             }
+            removeFinishedCallbacksLocked();
+        }
+    }
+
+    /**
+     * Removes any callbacks that notified us through {@link #onConsumerReadyToBeRemoved(Consumer)}
+     * that they are ready to be removed.
+     */
+    @GuardedBy("mLock")
+    private void removeFinishedCallbacksLocked() {
+        for (Consumer<T> callback: mCallbacksToRemove) {
+            mCallbacks.remove(callback);
+        }
+        mCallbacksToRemove.clear();
+    }
+
+    @Override
+    public void onConsumerReadyToBeRemoved(Consumer<T> callback) {
+        synchronized (mLock) {
+            mCallbacksToRemove.add(callback);
         }
     }
 }
\ No newline at end of file