com.twitter.common.net.loadbalancing.LeastConnectedStrategy.java Source code

Java tutorial

Introduction

Here is the source code for com.twitter.common.net.loadbalancing.LeastConnectedStrategy.java

Source

// =================================================================================================
// Copyright 2011 Twitter, Inc.
// -------------------------------------------------------------------------------------------------
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this work except in compliance with the License.
// You may obtain a copy of the License in the LICENSE file, or 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 com.twitter.common.net.loadbalancing;

import java.util.Collection;
import java.util.Map;
import java.util.Set;
import java.util.SortedSet;
import java.util.logging.Logger;

import com.google.common.base.Preconditions;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import com.google.common.collect.Sets;

import com.twitter.common.net.pool.ResourceExhaustedException;

/**
 * A load balancer that attempts to direct load towards a backend that has the fewest leased
 * connections.
 *
 * @author William Farner
 */
public class LeastConnectedStrategy<S> extends StaticLoadBalancingStrategy<S> {
    private static final Logger LOG = Logger.getLogger(LeastConnectedStrategy.class.getName());

    // Maps from backends to the number of connections made to them.
    private final Map<S, ConnectionStats> connections = Maps.newHashMap();

    // Manages sorting of connection counts, with a reference back to the backend.
    private final SortedSet<ConnectionStats> connectionStats = Sets.newTreeSet();

    /**
     * Encapsulates a set of connection stats that allow connections to be sorted as per the least
     * connected strategy.
     */
    private class ConnectionStats implements Comparable<ConnectionStats> {
        final S connectionKey;
        final int connectionId;
        int activeCount = 0; // Stores the total number of active connections.
        long useCount = 0; // Stores the total number times a connection has been used.

        ConnectionStats(S connectionKey, int connectionId) {
            this.connectionKey = connectionKey;
            this.connectionId = connectionId;
        }

        @Override
        public int compareTo(ConnectionStats other) {
            // Sort by number of active connections first.
            int difference = activeCount - other.activeCount;
            if (difference != 0) {
                return difference;
            }

            // Sub-sort by total number of times a connection has been used (this will ensure that
            // all backends are exercised).
            long useDifference = useCount - other.useCount;
            if (useDifference != 0) {
                return Long.signum(useDifference);
            }

            // If the above two are equal, break the tie using the connection id.
            return connectionId - other.connectionId;
        }

        @Override
        public boolean equals(Object o) {
            // We use ConnectionStats in a sorted container and so we need to have an equals
            // implementation consistent with compareTo, ie:
            // (x.compareTo(y) == 0) == x.equals(y)
            // We accomplish this directly.

            @SuppressWarnings("unchecked")
            ConnectionStats other = (ConnectionStats) o;
            return compareTo(other) == 0;
        }

        @Override
        public String toString() {
            return String.format("%d-%d", activeCount, useCount);
        }
    }

    @Override
    protected Collection<S> onBackendsOffered(Set<S> backends) {
        Map<S, ConnectionStats> newConnections = Maps.newHashMapWithExpectedSize(backends.size());
        Collection<ConnectionStats> newConnectionStats = Lists.newArrayListWithCapacity(backends.size());

        // Recreate all connection stats since their ordering may have changed and this is used for
        // comparison tie breaks.
        int backendId = 0;
        for (S backend : backends) {
            ConnectionStats stats = new ConnectionStats(backend, backendId++);

            // Retain the activeCount for existing backends to prevent dogpiling existing active servers
            ConnectionStats existing = connections.get(backend);
            if (existing != null) {
                stats.activeCount = existing.activeCount;
            }

            newConnections.put(backend, stats);
            newConnectionStats.add(stats);
        }

        connections.clear();
        connections.putAll(newConnections);
        connectionStats.clear();
        connectionStats.addAll(newConnectionStats);

        return connections.keySet();
    }

    @Override
    public S nextBackend() throws ResourceExhaustedException {
        Preconditions.checkState(connections.size() == connectionStats.size());

        if (connectionStats.isEmpty()) {
            throw new ResourceExhaustedException("No backends.");
        }

        return connectionStats.first().connectionKey;
    }

    @Override
    public void addConnectResult(S backendKey, ConnectionResult result, long connectTimeNanos) {
        Preconditions.checkNotNull(backendKey);
        Preconditions.checkState(connections.size() == connectionStats.size());
        Preconditions.checkNotNull(result);

        ConnectionStats stats = connections.get(backendKey);
        Preconditions.checkNotNull(stats);

        Preconditions.checkState(connectionStats.remove(stats));
        if (result == ConnectionResult.SUCCESS) {
            stats.activeCount++;
        }
        stats.useCount++;
        Preconditions.checkState(connectionStats.add(stats));
    }

    @Override
    public void connectionReturned(S backendKey) {
        Preconditions.checkNotNull(backendKey);
        Preconditions.checkState(connections.size() == connectionStats.size());

        ConnectionStats stats = connections.get(backendKey);
        Preconditions.checkNotNull(stats);

        if (stats.activeCount > 0) {
            Preconditions.checkState(connectionStats.remove(stats));
            stats.activeCount--;
            Preconditions.checkState(connectionStats.add(stats));
        } else {
            LOG.warning("connection stats dropped below zero, ignoring");
        }
    }
}