add groupBy overload with evictingMapFactory

davidmoten · davidmoten · commit 0099a1a63696 · 2016-05-12T17:31:22.000+10:00
diff --git a/src/main/java/rx/Observable.java b/src/main/java/rx/Observable.java
@@ -5885,6 +5885,60 @@ public final <K, R> Observable<GroupedObservable<K, R>> groupBy(final Func1<? su
         return lift(new OperatorGroupBy<T, K, R>(keySelector, elementSelector));
     }
     
+    /**
+     * Groups the items emitted by an {@code Observable} according to a specified criterion, and emits these
+     * grouped items as {@link GroupedObservable}s. The emitted {@code GroupedObservable} allows only a single 
+     * {@link Subscriber} during its lifetime and if this {@code Subscriber} unsubscribes before the 
+     * source terminates, the next emission by the source having the same key will trigger a new 
+     * {@code GroupedObservable} emission.
+     * <p>
+     * <img width="640" height="360" src="https://raw.github.com/wiki/ReactiveX/RxJava/images/rx-operators/groupBy.png" alt="">
+     * <p>
+     * <em>Note:</em> A {@link GroupedObservable} will cache the items it is to emit until such time as it
+     * is subscribed to. For this reason, in order to avoid memory leaks, you should not simply ignore those
+     * {@code GroupedObservable}s that do not concern you. Instead, you can signal to them that they may
+     * discard their buffers by applying an operator like {@link #ignoreElements} to them.
+     * <dl>
+     *  <dt><b>Scheduler:</b></dt>
+     *  <dd>{@code groupBy} does not operate by default on a particular {@link Scheduler}.</dd>
+     * </dl>
+     * 
+     * @param keySelector
+     *            a function that extracts the key for each item
+     * @param elementSelector
+     *            a function that extracts the return element for each item
+     * @param evictingMapFactory
+     *            a function that given an eviction action returns a {@link Map} instance that will be used to assign 
+     *            items to the appropriate {@code GroupedObservable}s. The {@code Map} instance must be thread-safe 
+     *            and any eviction must trigger a call to the supplied action (synchronously or asynchronously). 
+     *            This can be used to limit the size of the map by evicting keys by maximum size or access time for 
+     *            instance. If {@code evictingMapFactory} is null then no eviction strategy will be applied (and a suitable default thread-safe 
+     *            implementation of {@code Map} will be supplied). Here's an example using Guava's {@code CacheBuilder} from v19.0:
+     *            <pre>
+     *            {@code
+     *            Func1<Action1<K>, Map<K, Object>> mapFactory 
+     *              = action -> CacheBuilder.newBuilder()
+     *                  .maximumSize(1000)
+     *                  .expireAfterAccess(12, TimeUnit.HOUR)
+     *                  .removalListener(key -> action.call(key))
+     *                  .<K, Object> build().asMap();
+     *            }
+     *            </pre>
+     *            
+     * @param <K>
+     *            the key type
+     * @param <R>
+     *            the element type
+     * @return an {@code Observable} that emits {@link GroupedObservable}s, each of which corresponds to a
+     *         unique key value and each of which emits those items from the source Observable that share that
+     *         key value
+     * @see <a href="http://reactivex.io/documentation/operators/groupby.html">ReactiveX operators documentation: GroupBy</a>
+     */
+    @Experimental
+    public final <K, R> Observable<GroupedObservable<K, R>> groupBy(final Func1<? super T, ? extends K> keySelector, final Func1<? super T, ? extends R> elementSelector, final Func1<Action1<K>, Map<K, Object>> evictingMapFactory) {
+        return lift(new OperatorGroupBy<T, K, R>(keySelector, elementSelector, evictingMapFactory));
+    }
+    
     /**
      * Groups the items emitted by an {@code Observable} according to a specified criterion, and emits these
      * grouped items as {@link GroupedObservable}s. The emitted {@code GroupedObservable} allows only a single 
diff --git a/src/main/java/rx/internal/operators/OperatorGroupBy.java b/src/main/java/rx/internal/operators/OperatorGroupBy.java
@@ -21,10 +21,12 @@
 
 import rx.*;
 import rx.Observable.*;
+import rx.exceptions.Exceptions;
 import rx.functions.*;
 import rx.internal.producers.ProducerArbiter;
 import rx.internal.util.*;
 import rx.observables.GroupedObservable;
+import rx.observers.Subscribers;
 import rx.plugins.RxJavaPlugins;
 import rx.subscriptions.Subscriptions;
 
@@ -46,35 +48,50 @@ public final class OperatorGroupBy<T, K, V> implements Operator<GroupedObservabl
     final Func1<? super T, ? extends V> valueSelector;
     final int bufferSize;
     final boolean delayError;
+    final Func1<Action1<K>, Map<K, Object>> mapFactory; //nullable
     
     @SuppressWarnings({ "unchecked", "rawtypes" })
     public OperatorGroupBy(Func1<? super T, ? extends K> keySelector) {
-        this(keySelector, (Func1)UtilityFunctions.<T>identity(), RxRingBuffer.SIZE, false);
+        this(keySelector, (Func1)UtilityFunctions.<T>identity(), RxRingBuffer.SIZE, false, null);
     }
 
     public OperatorGroupBy(Func1<? super T, ? extends K> keySelector, Func1<? super T, ? extends V> valueSelector) {
-        this(keySelector, valueSelector, RxRingBuffer.SIZE, false);
+        this(keySelector, valueSelector, RxRingBuffer.SIZE, false, null);
+    }
+    
+    public OperatorGroupBy(Func1<? super T, ? extends K> keySelector, Func1<? super T, ? extends V> valueSelector, Func1<Action1<K>, Map<K, Object>> mapFactory) {
+        this(keySelector, valueSelector, RxRingBuffer.SIZE, false, mapFactory);
     }
 
-    public OperatorGroupBy(Func1<? super T, ? extends K> keySelector, Func1<? super T, ? extends V> valueSelector, int bufferSize, boolean delayError) {
+    public OperatorGroupBy(Func1<? super T, ? extends K> keySelector, Func1<? super T, ? extends V> valueSelector, int bufferSize, boolean delayError, Func1<Action1<K>, Map<K, Object>> mapFactory) {
         this.keySelector = keySelector;
         this.valueSelector = valueSelector;
         this.bufferSize = bufferSize;
         this.delayError = delayError;
+        this.mapFactory = mapFactory;
     }
     
     @Override
-    public Subscriber<? super T> call(Subscriber<? super GroupedObservable<K, V>> t) {
-        final GroupBySubscriber<T, K, V> parent = new GroupBySubscriber<T, K, V>(t, keySelector, valueSelector, bufferSize, delayError);
+    public Subscriber<? super T> call(Subscriber<? super GroupedObservable<K, V>> child) {
+        final GroupBySubscriber<T, K, V> parent;
+        try {
+            parent = new GroupBySubscriber<T, K, V>(child, keySelector, valueSelector, bufferSize, delayError, mapFactory);
+        } catch (Throwable ex) {
+            //Can reach here because mapFactory.call() may throw in constructor of GroupBySubscriber
+            Exceptions.throwOrReport(ex, child);
+            Subscriber<? super T> parent2 = Subscribers.empty();
+            parent2.unsubscribe();
+            return parent2;
+        }
 
-        t.add(Subscriptions.create(new Action0() {
+        child.add(Subscriptions.create(new Action0() {
             @Override
             public void call() {
                 parent.cancel();
             }
         }));
 
-        t.setProducer(parent.producer);
+        child.setProducer(parent.producer);
         
         return parent;
     }
@@ -101,6 +118,7 @@ public static final class GroupBySubscriber<T, K, V>
         final Map<Object, GroupedUnicast<K, V>> groups;
         final Queue<GroupedObservable<K, V>> queue;
         final GroupByProducer producer;
+        final Queue<K> evictedKeys;
         
         static final Object NULL_KEY = new Object();
         
@@ -129,18 +147,47 @@ public static final class GroupBySubscriber<T, K, V>
         static final AtomicIntegerFieldUpdater<GroupBySubscriber> WIP =
                 AtomicIntegerFieldUpdater.newUpdater(GroupBySubscriber.class, "wip");
         
-        public GroupBySubscriber(Subscriber<? super GroupedObservable<K, V>> actual, Func1<? super T, ? extends K> keySelector, Func1<? super T, ? extends V> valueSelector, int bufferSize, boolean delayError) {
+        public GroupBySubscriber(Subscriber<? super GroupedObservable<K, V>> actual, Func1<? super T, ? extends K> keySelector, 
+                Func1<? super T, ? extends V> valueSelector, int bufferSize, boolean delayError, 
+                Func1<Action1<K>, Map<K, Object>> mapFactory) {
             this.actual = actual;
             this.keySelector = keySelector;
             this.valueSelector = valueSelector;
             this.bufferSize = bufferSize;
             this.delayError = delayError;
-            this.groups = new ConcurrentHashMap<Object, GroupedUnicast<K, V>>();
             this.queue = new ConcurrentLinkedQueue<GroupedObservable<K, V>>();
             GROUP_COUNT.lazySet(this, 1);
             this.s = new ProducerArbiter();
             this.s.request(bufferSize);
             this.producer = new GroupByProducer(this);
+            if (mapFactory == null) {
+                this.groups = new ConcurrentHashMap<Object, GroupedUnicast<K, V>>();
+                this.evictedKeys = null;
+            } else {
+                this.evictedKeys = new ConcurrentLinkedQueue<K>();
+                this.groups = createMap(mapFactory, new EvictionAction<K>(evictedKeys));
+            }
+        }
+        
+        //declare a class instead of using anonymous class to 
+        //limit enclosing scope. Is this of value?
+        static class EvictionAction<K> implements Action1<K> {
+
+            final Queue<K> evictedKeys;
+
+            EvictionAction(Queue<K> evictedKeys) {
+                this.evictedKeys = evictedKeys;
+            }
+            
+            @Override
+            public void call(K key) {
+                evictedKeys.offer(key);
+            }
+        }
+        
+        @SuppressWarnings("unchecked")
+        private Map<Object, GroupedUnicast<K, V>> createMap(Func1<Action1<K>, Map<K, Object>> mapFactory, Action1<K> evictionAction) {
+            return (Map<Object, GroupedUnicast<K,V>>)(Map<Object, ?>) mapFactory.call(evictionAction);
         }
         
         @Override
@@ -196,6 +243,16 @@ public void onNext(T t) {
             }
 
             group.onNext(v);
+            
+            if (evictedKeys != null) {
+                K evictedKey;
+                while ((evictedKey = evictedKeys.poll()) != null) {
+                    GroupedUnicast<K, V> g = groups.get(evictedKey);
+                    if (g != null) {
+                        g.onComplete();
+                    }
+                }
+            }
 
             if (notNew) {
                 s.request(1);
@@ -224,6 +281,9 @@ public void onCompleted() {
                 e.onComplete();
             }
             groups.clear();
+            if (evictedKeys != null) {
+                evictedKeys.clear();
+            }
 
             done = true;
             GROUP_COUNT.decrementAndGet(this);
@@ -317,6 +377,9 @@ void errorAll(Subscriber<? super GroupedObservable<K, V>> a, Queue<?> q, Throwab
             q.clear();
             List<GroupedUnicast<K, V>> list = new ArrayList<GroupedUnicast<K, V>>(groups.values());
             groups.clear();
+            if (evictedKeys != null) { 
+                evictedKeys.clear();
+            }
             
             for (GroupedUnicast<K, V> e : list) {
                 e.onError(ex);
diff --git a/src/test/java/rx/internal/operators/OperatorGroupByTest.java b/src/test/java/rx/internal/operators/OperatorGroupByTest.java
@@ -1802,4 +1802,65 @@ public Integer call(Integer pair) {
         outer.assertValueCount(2);
 
     }
+    
+    @Test
+    public void mapFactoryEvictionWorks() {
+        Func1<Integer, Integer> keySelector = new Func1<Integer, Integer> (){
+            @Override
+            public Integer call(Integer t) {
+                return t /10;
+            }};
+        Func1<Integer, Integer> elementSelector = UtilityFunctions.identity();
+        final List<Integer> evictedKeys = new ArrayList<Integer>();
+        //normally we would use Guava CacheBuilder for instance but to save a test dependency 
+        //we make something custom
+        Func1<Action1<Integer>, Map<Integer, Object>> mapFactory = new Func1<Action1<Integer>, Map<Integer, Object>>() {
+            @Override
+            public Map<Integer, Object> call(final Action1<Integer> evicted) {
+                // is a bit risky to override the put method because 
+                // of possible side-effects (e.g. remove could call put and we did not know it)
+                // to fix just need to use composition but needs a verbose implementation of Map
+                // interface
+                return new ConcurrentHashMap<Integer,Object>() {
+                    private static final long serialVersionUID = -7519109652858021153L;
+                    
+                    Integer lastKey = null;
+                    
+                    @Override
+                    public Object put(Integer key, Object value) {
+                        if (this.size() >= 5) {
+                            super.remove(lastKey);
+                            evicted.call(lastKey);
+                            evictedKeys.add(lastKey);
+                        }
+                        Object result = super.put(key, value);
+                        lastKey = key;
+                        return result;
+                    }};
+            }};
+        TestSubscriber<String> ts = TestSubscriber.create();
+        Observable
+            .range(1, 100)
+            .groupBy(keySelector,elementSelector, mapFactory)
+            .flatMap(new Func1<GroupedObservable<Integer, Integer>, Observable<String>>() {
+                @Override
+                public Observable<String> call(final GroupedObservable<Integer, Integer> g) {
+                    return g.map(new Func1<Integer, String>() {
+                        @Override
+                        public String call(Integer x) {
+                            return g.getKey() + ":" + x;
+                        }
+                    });
+                }
+            })
+            .subscribe(ts);
+        assertEquals(Arrays.asList(4, 5, 6, 7, 8, 9), evictedKeys);
+        List<String> expected = Observable.range(1, 100).map(new Func1<Integer, String>() {
+            @Override
+            public String call(Integer x) {
+                return (x /10) + ":" + x;
+            }
+        }).toList().toBlocking().single();
+        assertEquals(expected, ts.getOnNextEvents());
+    }
 }
