AbstractLoadBalancer.java

package netflix.ocelli.rxnetty.internal;

import io.reactivex.netty.channel.Connection;
import io.reactivex.netty.client.ConnectionFactory;
import io.reactivex.netty.client.ConnectionObservable;
import io.reactivex.netty.client.ConnectionObservable.AbstractOnSubscribeFunc;
import io.reactivex.netty.client.ConnectionProvider;
import io.reactivex.netty.client.pool.PooledConnectionProvider;
import io.reactivex.netty.protocol.tcp.client.events.TcpClientEventListener;
import netflix.ocelli.Instance;
import netflix.ocelli.LoadBalancerStrategy;
import netflix.ocelli.rxnetty.FailureListener;
import rx.Observable;
import rx.Subscriber;
import rx.functions.Action1;
import rx.functions.Func1;

import java.net.SocketAddress;
import java.util.List;
import java.util.NoSuchElementException;

/**
 * An abstract load balancer for all TCP based protocols.
 *
 * <h2>Failure detection</h2>
 *
 * For every host that this load balancer connects, it provides a way to register a {@link TcpClientEventListener}
 * instance that can detect failures based on the various events received. Upon detecting the failure, an appropriate
 * action can be taken for the host, using the provided {@link FailureListener}.
 *
 * <h2>Use with RxNetty clients</h2>
 *
 * In order to use this load balancer with RxNetty clients, one has to convert it to an instance of
 * {@link ConnectionProvider} by calling {@link #toConnectionProvider()}
 *
 * @param <W> Type of Objects written on the connections created by this load balancer.
 * @param <R> Type of Objects read from the connections created by this load balancer.
 */
public abstract class AbstractLoadBalancer<W, R> {

    protected final Observable<Instance<SocketAddress>> hosts;
    protected final LoadBalancerStrategy<HostConnectionProvider<W, R>> loadBalancer;
    protected final Func1<FailureListener, ? extends TcpClientEventListener> eventListenerFactory;

    protected AbstractLoadBalancer(Observable<Instance<SocketAddress>> hosts,
                                   Func1<FailureListener, ? extends TcpClientEventListener> eventListenerFactory,
                                   LoadBalancerStrategy<HostConnectionProvider<W, R>> loadBalancer) {
        this.hosts = hosts;
        this.eventListenerFactory = eventListenerFactory;
        this.loadBalancer = loadBalancer;
    }

    /**
     * Converts this load balancer to a {@link ConnectionProvider} to be used with RxNetty clients.
     *
     * @return {@link ConnectionProvider} for this load balancer.
     */
    public ConnectionProvider<W, R> toConnectionProvider() {
        return ConnectionProvider.create(new Func1<ConnectionFactory<W, R>, ConnectionProvider<W, R>>() {
            @Override
            public ConnectionProvider<W, R> call(final ConnectionFactory<W, R> connectionFactory) {
                return toConnectionProvider(connectionFactory);
            }
        });
    }

    /*Visible for testing*/ ConnectionProvider<W, R> toConnectionProvider(final ConnectionFactory<W, R> factory) {

        final Observable<Instance<ConnectionProvider<W, R>>> providerStream =
                hosts.map(new Func1<Instance<SocketAddress>, Instance<ConnectionProvider<W, R>>>() {
                    @Override
                    public Instance<ConnectionProvider<W, R>> call(final Instance<SocketAddress> host) {
                        final ConnectionProvider<W, R> pcp = newConnectionProviderForHost(host, factory);

                        return new Instance<ConnectionProvider<W, R>>() {
                            @Override
                            public Observable<Void> getLifecycle() {
                                return host.getLifecycle();
                            }

                            @Override
                            public ConnectionProvider<W, R> getValue() {
                                return pcp;
                            }
                        };
                    }
                });

        return new LoadBalancingProvider(factory, providerStream);
    }

    protected ConnectionProvider<W, R> newConnectionProviderForHost(Instance<SocketAddress> host,
                                                                    ConnectionFactory<W, R> connectionFactory) {
        /*
         * Bounds on the concurrency (concurrent connections) should be enforced at the request
         * processing level, providing a bound on number of connections is a difficult number
         * to determine.
         */
        return PooledConnectionProvider.createUnbounded(connectionFactory, host.getValue());
    }

    /*Visible for testing*/class LoadBalancingProvider extends ConnectionProvider<W, R> {

        private final HostHolder<W, R> hostHolder;

        public LoadBalancingProvider(ConnectionFactory<W, R> connectionFactory,
                                     Observable<Instance<ConnectionProvider<W, R>>> providerStream) {
            super(connectionFactory);
            hostHolder = new HostHolder<>(providerStream, eventListenerFactory);
        }

        @Override
        public ConnectionObservable<R, W> nextConnection() {

            return ConnectionObservable.createNew(new AbstractOnSubscribeFunc<R, W>() {
                @Override
                protected void doSubscribe(Subscriber<? super Connection<R, W>> sub,
                                           Action1<ConnectionObservable<R, W>> subscribeAllListenersAction) {
                    final List<HostConnectionProvider<W, R>> providers = hostHolder.getProviders();

                    if (null == providers || providers.isEmpty()) {
                        sub.onError(new NoSuchElementException("No hosts available."));
                    }

                    HostConnectionProvider<W, R> hcp = loadBalancer.choose(providers);

                    ConnectionObservable<R, W> nextConnection = hcp.getProvider().nextConnection();
                    if (hcp.getEventsListener() != null) {
                        nextConnection.subscribeForEvents(hcp.getEventsListener());
                    }
                    subscribeAllListenersAction.call(nextConnection);
                    nextConnection.unsafeSubscribe(sub);
                }
            });
        }

        @Override
        protected Observable<Void> doShutdown() {
            return hostHolder.shutdown();
        }
    }
}