netflix.ocelli.loadbalancer.DefaultPartitioningLoadBalancer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ocelli-core Show documentation
ocelli-core developed by Netflix
There is a newer version: 0.1.0-rc.2
package netflix.ocelli.loadbalancer;

import netflix.ocelli.ClientConnector;
import netflix.ocelli.FailureDetectorFactory;
import netflix.ocelli.LoadBalancer;
import netflix.ocelli.LoadBalancers;
import netflix.ocelli.MembershipEvent;
import netflix.ocelli.PartitionedLoadBalancer;
import netflix.ocelli.WeightingStrategy;
import netflix.ocelli.selectors.ClientsAndWeights;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import rx.Observable;
import rx.functions.Action1;
import rx.functions.Func1;
import rx.subjects.PublishSubject;
import rx.subscriptions.CompositeSubscription;

import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.ConcurrentMap;

public class DefaultPartitioningLoadBalancer implements PartitionedLoadBalancer {
    private static final Logger LOG = LoggerFactory.getLogger(DefaultPartitioningLoadBalancer.class);
    
    private final CompositeSubscription cs = new CompositeSubscription();
    private final Func1> partitioner;
    private final Observable> hostSource;
    private final ConcurrentMap partitions = new ConcurrentHashMap();
    private final WeightingStrategy weightingStrategy;
    private final FailureDetectorFactory failureDetector;
    private final ClientConnector clientConnector;
    private final Func1 connectedHostCountStrategy;
    private final Func1 quaratineDelayStrategy;
    private final Func1, Observable> selectionStrategy;
    private final String name;
    
    private final class Holder {
        final PublishSubject> hostStream;
        final LoadBalancer loadBalancer;
        
        public Holder(LoadBalancer loadBalancer, PublishSubject> hostStream) {
            this.loadBalancer = loadBalancer;
            this.hostStream = hostStream;
        }
    }
    
    public DefaultPartitioningLoadBalancer(
            String name,
            Observable> hostSource,
            ClientConnector clientConnector,
            FailureDetectorFactory failureDetector,
            Func1, Observable> selectionStrategy,
            Func1 quaratineDelayStrategy,
            Func1 connectedHostCountStrategy,
            WeightingStrategy weightingStrategy,
            Func1> partitioner) {
        
        this.partitioner            = partitioner;
        this.hostSource             = hostSource;
        this.failureDetector        = failureDetector;
        this.clientConnector        = clientConnector;
        this.selectionStrategy      = selectionStrategy;
        this.weightingStrategy      = weightingStrategy;
        this.quaratineDelayStrategy = quaratineDelayStrategy;
        this.name                   = name;
        this.connectedHostCountStrategy = connectedHostCountStrategy;
        
        initialize();
    }

    private void initialize() {
        cs.add(hostSource
            .subscribe(new Action1>() {
                @Override
                public void call(final MembershipEvent event) {
                    partitioner
                            .call(event.getClient())
                            .subscribe(new Action1() {
                                @Override
                                public void call(K id) {
                                    getOrCreateHolder(id).hostStream.onNext(event);
                                }
                            });
                }
            })
        );
    }
    
    @Override
    public void shutdown() {
        cs.unsubscribe();
    }
    
    private Holder getOrCreateHolder(K id) {
        Holder holder = partitions.get(id);
        if (null == holder) {
            PublishSubject> subject = PublishSubject.create();
            Holder newHolder = new Holder(createPartition(id, subject), subject);
            holder = partitions.putIfAbsent(id, newHolder);
            if (holder == null) {
                holder = newHolder;
            }
        }
        return holder;
    }
    
    @Override
    public LoadBalancer get(K id) {
        return getOrCreateHolder(id).loadBalancer;
    }

    @Override
    public Observable listKeys() {
        return Observable.from(partitions.keySet());
    }
    
    private LoadBalancer createPartition(K id, Observable> hostSource) {
        LOG.info("Creating partition : " + id);
        return LoadBalancers.newBuilder(hostSource)
                .withName(getName() + "_" + id)
                .withQuarantineStrategy(quaratineDelayStrategy)
                .withSelectionStrategy(selectionStrategy)
                .withWeightingStrategy(weightingStrategy)
                .withActiveClientCountStrategy(connectedHostCountStrategy)
                .withClientConnector(clientConnector)
                .withFailureDetector(failureDetector)
                .build();
    }

    private String getName() {
        return this.name;
    }
}