io.reactivx.mantis.operators.OperatorGroupBy Maven / Gradle / Ivy
/*
* Copyright 2019 Netflix, Inc.
*
* Licensed under the Apache License, Version 2.0 (the "License");
* you may not use this file except in compliance with the License.
* You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing, software
* distributed under the License is distributed on an "AS IS" BASIS,
* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
* See the License for the specific language governing permissions and
* limitations under the License.
*/
package io.reactivx.mantis.operators;
import java.util.Queue;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentLinkedQueue;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import java.util.concurrent.atomic.AtomicLong;
import java.util.concurrent.atomic.AtomicLongFieldUpdater;
import rx.Observable;
import rx.Observable.OnSubscribe;
import rx.Observable.Operator;
import rx.Observer;
import rx.Producer;
import rx.Subscriber;
import rx.exceptions.OnErrorThrowable;
import rx.functions.Action0;
import rx.functions.Func1;
import rx.internal.operators.NotificationLite;
import rx.observables.GroupedObservable;
import rx.subjects.Subject;
/**
* Groups the items emitted by an Observable according to a specified criterion, and emits these
* grouped items as Observables, one Observable per group.
*
* 
*
* @param
* the key type
* @param
* the source and group value type
* @param
* the value type of the groups
*/
public class OperatorGroupBy implements Operator, T> {
private final static Func1 IDENTITY = new Func1() {
@Override
public Object call(Object t) {
return t;
}
};
final Func1 keySelector;
final Func1 valueSelector;
@SuppressWarnings("unchecked")
public OperatorGroupBy(final Func1 keySelector) {
this(keySelector, (Func1) IDENTITY);
}
public OperatorGroupBy(
Func1 keySelector,
Func1 valueSelector) {
this.keySelector = keySelector;
this.valueSelector = valueSelector;
}
@Override
public Subscriber call(final Subscriber> child) {
return new GroupBySubscriber(keySelector, valueSelector, child);
}
static final class GroupBySubscriber extends Subscriber {
@SuppressWarnings("rawtypes")
static final AtomicIntegerFieldUpdater COMPLETION_EMITTED_UPDATER = AtomicIntegerFieldUpdater.newUpdater(GroupBySubscriber.class, "completionEmitted");
@SuppressWarnings("rawtypes")
static final AtomicIntegerFieldUpdater TERMINATED_UPDATER = AtomicIntegerFieldUpdater.newUpdater(GroupBySubscriber.class, "terminated");
@SuppressWarnings("rawtypes")
static final AtomicLongFieldUpdater REQUESTED = AtomicLongFieldUpdater.newUpdater(GroupBySubscriber.class, "requested");
@SuppressWarnings("rawtypes")
static final AtomicLongFieldUpdater BUFFERED_COUNT = AtomicLongFieldUpdater.newUpdater(GroupBySubscriber.class, "bufferedCount");
private static final int MAX_QUEUE_SIZE = 1024;
final GroupBySubscriber self = this;
final Func1 keySelector;
final Func1 elementSelector;
final Subscriber> child;
private final ConcurrentHashMap> groups = new ConcurrentHashMap>();
volatile int completionEmitted;
volatile int terminated;
volatile long requested;
volatile long bufferedCount;
public GroupBySubscriber(
Func1 keySelector,
Func1 elementSelector,
Subscriber> child) {
super();
this.keySelector = keySelector;
this.elementSelector = elementSelector;
this.child = child;
}
@Override
public void onStart() {
REQUESTED.set(this, MAX_QUEUE_SIZE);
request(MAX_QUEUE_SIZE);
}
@Override
public void onCompleted() {
if (TERMINATED_UPDATER.compareAndSet(this, 0, 1)) {
// if we receive onCompleted from our parent we onComplete children
// for each group check if it is ready to accept more events if so pass the oncomplete through else buffer it.
for (GroupState group : groups.values()) {
emitItem(group, NotificationLite.completed());
}
// special case (no groups emitted ... or all unsubscribed)
if (groups.size() == 0) {
// we must track 'completionEmitted' seperately from 'completed' since `completeInner` can result in childObserver.onCompleted() being emitted
if (COMPLETION_EMITTED_UPDATER.compareAndSet(this, 0, 1)) {
child.onCompleted();
}
}
}
}
@Override
public void onError(Throwable e) {
if (TERMINATED_UPDATER.compareAndSet(this, 0, 1)) {
// we immediately tear everything down if we receive an error
child.onError(e);
}
}
void requestFromGroupedObservable(long n, GroupState group) {
group.requested.getAndAdd(n);
if (group.count.getAndIncrement() == 0) {
pollQueue(group);
}
}
// The grouped observable propagates the 'producer.request' call from it's subscriber to this method
// Here we keep track of the requested count for each group
// If we already have items queued when a request comes in we vend those and decrement the outstanding request count
@Override
public void onNext(T t) {
try {
final K key = keySelector.call(t);
GroupState group = groups.get(key);
if (group == null) {
// this group doesn't exist
if (child.isUnsubscribed()) {
// we have been unsubscribed on the outer so won't send any more groups
return;
}
group = createNewGroup(key);
}
emitItem(group, NotificationLite.next(t));
} catch (Throwable e) {
onError(OnErrorThrowable.addValueAsLastCause(e, t));
}
}
private GroupState createNewGroup(final K key) {
final GroupState groupState = new GroupState();
GroupedObservable go = GroupedObservable.create(key, new OnSubscribe() {
@Override
public void call(final Subscriber o) {
o.setProducer(new Producer() {
@Override
public void request(long n) {
requestFromGroupedObservable(n, groupState);
}
});
final AtomicBoolean once = new AtomicBoolean();
groupState.getObservable().doOnUnsubscribe(new Action0() {
@Override
public void call() {
if (once.compareAndSet(false, true)) {
// done once per instance, either onComplete or onUnSubscribe
cleanupGroup(key);
}
}
}).unsafeSubscribe(new Subscriber(o) {
@Override
public void onCompleted() {
o.onCompleted();
// eagerly cleanup instead of waiting for unsubscribe
if (once.compareAndSet(false, true)) {
// done once per instance, either onComplete or onUnSubscribe
cleanupGroup(key);
}
}
@Override
public void onError(Throwable e) {
o.onError(e);
}
@Override
public void onNext(T t) {
try {
o.onNext(elementSelector.call(t));
} catch (Throwable e) {
onError(OnErrorThrowable.addValueAsLastCause(e, t));
}
}
});
}
});
GroupState putIfAbsent = groups.putIfAbsent(key, groupState);
if (putIfAbsent != null) {
// this shouldn't happen (because we receive onNext sequentially) and would mean we have a bug
throw new IllegalStateException("Group already existed while creating a new one");
}
child.onNext(go);
return groupState;
}
private void cleanupGroup(K key) {
GroupState removed;
removed = groups.remove(key);
if (removed != null) {
if (removed.buffer.size() > 0) {
BUFFERED_COUNT.addAndGet(self, -removed.buffer.size());
}
completeInner();
// since we may have unsubscribed early with items in the buffer
// we remove those above and have freed up room to request more
// so give it a chance to request more now
requestMoreIfNecessary();
}
}
private void emitItem(GroupState groupState, Object item) {
Queue