All Downloads are FREE. Search and download functionalities are using the official Maven repository.

rx.internal.operators.OperatorGroupBy Maven / Gradle / Ivy

Go to download 

/**
 * Copyright 2014 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 rx.internal.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.exceptions.*;
import rx.Observer;
import rx.Producer;
import rx.Subscriber;
import rx.functions.Action0;
import rx.functions.Func1;
import rx.observables.GroupedObservable;
import rx.subjects.Subject;
import rx.subscriptions.Subscriptions;

/**
 * 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> {
    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 {
        private static final int MAX_QUEUE_SIZE = 1024;
        final GroupBySubscriber self = this;
        final Func1 keySelector;
        final Func1 elementSelector;
        final Subscriber> child;

        // We should not call `unsubscribe()` until `groups.isEmpty() && child.isUnsubscribed()` is true.
        // Use `WIP_FOR_UNSUBSCRIBE_UPDATER` to monitor these statuses and call `unsubscribe()` properly.
        // Should check both when `child.unsubscribe` is called and any group is removed.
        @SuppressWarnings("rawtypes")
        static final AtomicIntegerFieldUpdater WIP_FOR_UNSUBSCRIBE_UPDATER = AtomicIntegerFieldUpdater.newUpdater(GroupBySubscriber.class, "wipForUnsubscribe");
        volatile int wipForUnsubscribe = 1;

        public GroupBySubscriber(
                Func1 keySelector,
                Func1 elementSelector,
                Subscriber> child) {
            super();
            this.keySelector = keySelector;
            this.elementSelector = elementSelector;
            this.child = child;
            child.add(Subscriptions.create(new Action0() {

                @Override
                public void call() {
                    if (WIP_FOR_UNSUBSCRIBE_UPDATER.decrementAndGet(self) == 0) {
                        self.unsubscribe();
                    }
                }

            }));
        }

        private static class GroupState {
            private final Subject s = BufferUntilSubscriber.create();
            private final AtomicLong requested = new AtomicLong();
            private final AtomicLong count = new AtomicLong();
            private final Queue buffer = new ConcurrentLinkedQueue(); // TODO should this be lazily created?

            public Observable getObservable() {
                return s;
            }

            public Observer getObserver() {
                return s;
            }

        }

        private final ConcurrentHashMap> groups = new ConcurrentHashMap>();

        private static final NotificationLite nl = NotificationLite.instance();

        volatile int completionEmitted;

        private static final int UNTERMINATED = 0;
        private static final int TERMINATED_WITH_COMPLETED = 1;
        private static final int TERMINATED_WITH_ERROR = 2;

        // Must be one of `UNTERMINATED`, `TERMINATED_WITH_COMPLETED`, `TERMINATED_WITH_ERROR`
        volatile int terminated = UNTERMINATED;

        @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");

        volatile long requested;
        @SuppressWarnings("rawtypes")
        static final AtomicLongFieldUpdater REQUESTED = AtomicLongFieldUpdater.newUpdater(GroupBySubscriber.class, "requested");

        volatile long bufferedCount;
        @SuppressWarnings("rawtypes")
        static final AtomicLongFieldUpdater BUFFERED_COUNT = AtomicLongFieldUpdater.newUpdater(GroupBySubscriber.class, "bufferedCount");

        @Override
        public void onStart() {
            REQUESTED.set(this, MAX_QUEUE_SIZE);
            request(MAX_QUEUE_SIZE);
        }

        @Override
        public void onCompleted() {
            if (TERMINATED_UPDATER.compareAndSet(this, UNTERMINATED, TERMINATED_WITH_COMPLETED)) {
                // 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, nl.completed());
                }

                // special case (no groups emitted ... or all unsubscribed)
                if (groups.isEmpty()) {
                    // 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, UNTERMINATED, TERMINATED_WITH_ERROR)) {
                // It's safe to access all groups and emit the error.
                // onNext and onError are in sequence so no group will be created in the loop.
                for (GroupState group : groups.values()) {
                    emitItem(group, nl.error(e));
                }
                try {
                    // we immediately tear everything down if we receive an error
                    child.onError(e);
                } finally {
                    // We have not chained the subscribers, so need to call it explicitly.
                    unsubscribe();
                }
            }
        }

        // 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

        void requestFromGroupedObservable(long n, GroupState group) {
            BackpressureUtils.getAndAddRequest(group.requested, n);
            if (group.count.getAndIncrement() == 0) {
                pollQueue(group);
            }
        }

        private Object groupedKey(K key) {
            return key == null ? NULL_KEY : key;
        }

        @SuppressWarnings("unchecked")
        private K getKey(Object groupedKey) {
            return groupedKey == NULL_KEY ? null : (K) groupedKey;
        }

        @Override
        public void onNext(T t) {
            try {
                final Object key = groupedKey(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);
                }
                if (group != null) {
                    emitItem(group, nl.next(t));
                }
            } catch (Throwable e) {
                Exceptions.throwOrReport(e, this, t);
            }
        }

        private GroupState createNewGroup(final Object key) {
            final GroupState groupState = new GroupState();

            GroupedObservable go = GroupedObservable.create(getKey(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);
                            // 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 onNext(T t) {
                            try {
                                o.onNext(elementSelector.call(t));
                            } catch (Throwable e) {
                                Exceptions.throwOrReport(e, this, t);
                            }
                        }
                    });
                }
            });

            GroupState putIfAbsent;
            for (;;) {
                int wip = wipForUnsubscribe;
                if (wip <= 0) {
                    return null;
                }
                if (WIP_FOR_UNSUBSCRIBE_UPDATER.compareAndSet(this, wip, wip + 1)) {
                    putIfAbsent = groups.putIfAbsent(key, groupState);
                    break;
                }
            }
            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(Object key) {
            GroupState removed;
            removed = groups.remove(key);
            if (removed != null) {
                if (!removed.buffer.isEmpty()) {
                    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 q = groupState.buffer;
            AtomicLong keyRequested = groupState.requested;
            //don't need to check for requested being Long.MAX_VALUE because this
            //field is capped at MAX_QUEUE_SIZE
            REQUESTED.decrementAndGet(this);
            // short circuit buffering
            if (keyRequested != null && keyRequested.get() > 0 && (q == null || q.isEmpty())) {
                @SuppressWarnings("unchecked")
                Observer obs = (Observer)groupState.getObserver();
                nl.accept(obs, item);
                if (keyRequested.get() != Long.MAX_VALUE) {
                    // best endeavours check (no CAS loop here) because we mainly care about 
                    // the initial request being Long.MAX_VALUE and that value being conserved.
                    keyRequested.decrementAndGet();
                }
            } else {
                q.add(item);
                BUFFERED_COUNT.incrementAndGet(this);

                if (groupState.count.getAndIncrement() == 0) {
                    pollQueue(groupState);
                }
            }
            requestMoreIfNecessary();
        }

        private void pollQueue(GroupState groupState) {
            do {
                drainIfPossible(groupState);
                long c = groupState.count.decrementAndGet();
                if (c > 1) {

                    /*
                     * Set down to 1 and then iterate again.
                     * we lower it to 1 otherwise it could have grown very large while in the last poll loop
                     * and then we can end up looping all those times again here before existing even once we've drained
                     */
                    groupState.count.set(1);
                    // we now loop again, and if anything tries scheduling again after this it will increment and cause us to loop again after
                }
            } while (groupState.count.get() > 0);
        }

        private void requestMoreIfNecessary() {
            if (REQUESTED.get(this) == 0 && terminated == 0) {
                long toRequest = MAX_QUEUE_SIZE - BUFFERED_COUNT.get(this);
                if (toRequest > 0 && REQUESTED.compareAndSet(this, 0, toRequest)) {
                    request(toRequest);
                }
            }
        }

        private void drainIfPossible(GroupState groupState) {
            while (groupState.requested.get() > 0) {
                Object t = groupState.buffer.poll();
                if (t != null) {
                    @SuppressWarnings("unchecked")
                    Observer obs = (Observer)groupState.getObserver();
                    nl.accept(obs, t);
                    if (groupState.requested.get()!=Long.MAX_VALUE) {
                        // best endeavours check (no CAS loop here) because we mainly care about 
                        // the initial request being Long.MAX_VALUE and that value being conserved.
                        groupState.requested.decrementAndGet();
                    }
                    BUFFERED_COUNT.decrementAndGet(this);

                    // if we have used up all the events we requested from upstream then figure out what to ask for this time based on the empty space in the buffer
                    requestMoreIfNecessary();
                } else {
                    // queue is empty break
                    break;
                }
            }
        }

        private void completeInner() {
            // A group is removed, so check if we need to call `unsubscribe`
            if (WIP_FOR_UNSUBSCRIBE_UPDATER.decrementAndGet(this) == 0) {
                // It means `groups.isEmpty() && child.isUnsubscribed()` is true
                unsubscribe();
            } else if (groups.isEmpty() && terminated == TERMINATED_WITH_COMPLETED) {
                // if we have no outstanding groups (all completed or unsubscribe) and terminated on outer
                // completionEmitted ensures we only emit onCompleted once
                if (COMPLETION_EMITTED_UPDATER.compareAndSet(this, 0, 1)) {
                    child.onCompleted();
                }
            }
        }

    }

    private final static Func1 IDENTITY = new Func1() {
        @Override
        public Object call(Object t) {
            return t;
        }
    };

    private static final Object NULL_KEY = new Object();
}
 















© 2015 - 2024 Weber Informatics LLC | Privacy Policy