Fix NPE when the key is null in GroupBy

zsxwing · zsxwing · commit d5b4d6f380a0 · 2014-12-12T23:15:44.000+08:00
diff --git a/src/main/java/rx/internal/operators/OperatorGroupBy.java b/src/main/java/rx/internal/operators/OperatorGroupBy.java
@@ -102,7 +102,7 @@ public Observer<T> getObserver() {
 
         }
 
-        private final ConcurrentHashMap<K, GroupState<K, T>> groups = new ConcurrentHashMap<K, GroupState<K, T>>();
+        private final ConcurrentHashMap<Object, GroupState<K, T>> groups = new ConcurrentHashMap<Object, GroupState<K, T>>();
 
         private static final NotificationLite<Object> nl = NotificationLite.instance();
 
@@ -166,10 +166,18 @@ void requestFromGroupedObservable(long n, GroupState<K, T> group) {
             }
         }
 
+        private Object groupedKey(K key) {
+            return key == null ? NULL_KEY : key;
+        }
+
+        private K getKey(Object groupedKey) {
+            return groupedKey == NULL_KEY ? null : (K) groupedKey;
+        }
+
         @Override
         public void onNext(T t) {
             try {
-                final K key = keySelector.call(t);
+                final Object key = groupedKey(keySelector.call(t));
                 GroupState<K, T> group = groups.get(key);
                 if (group == null) {
                     // this group doesn't exist
@@ -185,10 +193,10 @@ public void onNext(T t) {
             }
         }
 
-        private GroupState<K, T> createNewGroup(final K key) {
+        private GroupState<K, T> createNewGroup(final Object key) {
             final GroupState<K, T> groupState = new GroupState<K, T>();
 
-            GroupedObservable<K, R> go = GroupedObservable.create(key, new OnSubscribe<R>() {
+            GroupedObservable<K, R> go = GroupedObservable.create(getKey(key), new OnSubscribe<R>() {
 
                 @Override
                 public void call(final Subscriber<? super R> o) {
@@ -252,7 +260,7 @@ public void onNext(T t) {
             return groupState;
         }
 
-        private void cleanupGroup(K key) {
+        private void cleanupGroup(Object key) {
             GroupState<K, T> removed;
             removed = groups.remove(key);
             if (removed != null) {
@@ -357,4 +365,5 @@ public Object call(Object t) {
         }
     };
 
+    private static final Object NULL_KEY = new Object();
 }
diff --git a/src/test/java/rx/internal/operators/OperatorGroupByTest.java b/src/test/java/rx/internal/operators/OperatorGroupByTest.java
@@ -28,6 +28,7 @@
 import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
+import java.util.List;
 import java.util.Map;
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentLinkedQueue;
@@ -1357,4 +1358,31 @@ public Observable<Integer> call(GroupedObservable<Integer, Integer> t) {
 
     };
 
-}
+    @Test
+    public void testGroupByWithNullKey() {
+        final String[] key = new String[]{"uninitialized"};
+        final List<String> values = new ArrayList<String>();
+        Observable.just("a", "b", "c").groupBy(new Func1<String, String>() {
+
+            @Override
+            public String call(String value) {
+                return null;
+            }
+        }).subscribe(new Action1<GroupedObservable<String, String>>() {
+
+            @Override
+            public void call(GroupedObservable<String, String> groupedObservable) {
+                key[0] = groupedObservable.getKey();
+                groupedObservable.subscribe(new Action1<String>() {
+
+                    @Override
+                    public void call(String s) {
+                        values.add(s);
+                    }
+                });
+            }
+        });
+        assertEquals(null, key[0]);
+        assertEquals(Arrays.asList("a", "b", "c"), values);
+    }
+}

Original file line number	Diff line number	Diff line change
`@@ -102,7 +102,7 @@ public Observer<T> getObserver() {`
`102`	`102`
`103`	`103`	`}`
`104`	`104`
`105`		`- private final ConcurrentHashMap<K, GroupState<K, T>> groups = new ConcurrentHashMap<K, GroupState<K, T>>();`
	`105`	`+ private final ConcurrentHashMap<Object, GroupState<K, T>> groups = new ConcurrentHashMap<Object, GroupState<K, T>>();`
`106`	`106`
`107`	`107`	`private static final NotificationLite<Object> nl = NotificationLite.instance();`
`108`	`108`
`@@ -166,10 +166,18 @@ void requestFromGroupedObservable(long n, GroupState<K, T> group) {`
`166`	`166`	`}`
`167`	`167`	`}`
`168`	`168`
	`169`	`+ private Object groupedKey(K key) {`
	`170`	`+ return key == null ? NULL_KEY : key;`
	`171`	`+ }`
	`172`	`+`
	`173`	`+ private K getKey(Object groupedKey) {`
	`174`	`+ return groupedKey == NULL_KEY ? null : (K) groupedKey;`
	`175`	`+ }`
	`176`	`+`
`169`	`177`	`@Override`
`170`	`178`	`public void onNext(T t) {`
`171`	`179`	`try {`
`172`		`- final K key = keySelector.call(t);`
	`180`	`+ final Object key = groupedKey(keySelector.call(t));`
`173`	`181`	`GroupState<K, T> group = groups.get(key);`
`174`	`182`	`if (group == null) {`
`175`	`183`	`// this group doesn't exist`
`@@ -185,10 +193,10 @@ public void onNext(T t) {`
`185`	`193`	`}`
`186`	`194`	`}`
`187`	`195`
`188`		`- private GroupState<K, T> createNewGroup(final K key) {`
	`196`	`+ private GroupState<K, T> createNewGroup(final Object key) {`
`189`	`197`	`final GroupState<K, T> groupState = new GroupState<K, T>();`
`190`	`198`
`191`		`- GroupedObservable<K, R> go = GroupedObservable.create(key, new OnSubscribe<R>() {`
	`199`	`+ GroupedObservable<K, R> go = GroupedObservable.create(getKey(key), new OnSubscribe<R>() {`
`192`	`200`
`193`	`201`	`@Override`
`194`	`202`	`public void call(final Subscriber<? super R> o) {`
`@@ -252,7 +260,7 @@ public void onNext(T t) {`
`252`	`260`	`return groupState;`
`253`	`261`	`}`
`254`	`262`
`255`		`- private void cleanupGroup(K key) {`
	`263`	`+ private void cleanupGroup(Object key) {`
`256`	`264`	`GroupState<K, T> removed;`
`257`	`265`	`removed = groups.remove(key);`
`258`	`266`	`if (removed != null) {`
`@@ -357,4 +365,5 @@ public Object call(Object t) {`
`357`	`365`	`}`
`358`	`366`	`};`
`359`	`367`
	`368`	`+ private static final Object NULL_KEY = new Object();`
`360`	`369`	`}`