net.sf.hajdbc.balancer.load.LoadBalancer Maven / Gradle / Ivy

Go to download
Show more of this group Show more artifacts with this name
Show all versions of ha-jdbc2 Show documentation
High-Availability JDBC
There is a newer version: 3.6.61
/*
 * HA-JDBC: High-Availability JDBC
 * Copyright (C) 2012  Paul Ferraro
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Lesser General Public License as published by
 * the Free Software Foundation, either version 3 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General Public License
 * along with this program.  If not, see .
 */
package net.sf.hajdbc.balancer.load;

import java.util.*;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.concurrent.locks.Lock;
import java.util.concurrent.locks.ReentrantLock;

import net.sf.hajdbc.Database;
import net.sf.hajdbc.balancer.AbstractBalancer;
import net.sf.hajdbc.invocation.Invoker;
import net.sf.hajdbc.state.StateManager;
import net.sf.hajdbc.util.Collections;

/**
 * Balancer implementation whose {@link #next()} implementation returns the database with the least load.
 *
 * @author  Paul Ferraro
 * @param  either java.sql.Driver or javax.sql.DataSource
 */
public class LoadBalancer> extends AbstractBalancer
{
	private final Lock lock = new ReentrantLock();
	
	private volatile SortedMap databaseMap = Collections.emptySortedMap();

	private Comparator> comparator = new Comparator>()
	{
		@Override
		public int compare(Map.Entry mapEntry1, Map.Entry mapEntry2)
		{
			D database1 = mapEntry1.getKey();
			D database2 = mapEntry2.getKey();

			float load1 = mapEntry1.getValue().get();
			float load2 = mapEntry2.getValue().get();
			
			int weight1 = database1.getWeight();
			int weight2 = database2.getWeight();
			
			// If weights are the same, we can simply compare the loads
			if (weight1 == weight2)
			{
				return Float.compare(load1, load2);
			}
			
			float weightedLoad1 = (weight1 != 0) ? (load1 / weight1) : Float.POSITIVE_INFINITY;
			float weightedLoad2 = (weight2 != 0) ? (load2 / weight2) : Float.POSITIVE_INFINITY;
			
			return Float.compare(weightedLoad1, weightedLoad2);
		}
	};

	/**
	 * Constructs a new LoadBalancer
	 * @param databases
	 */
	public LoadBalancer(Set databases)
	{
		if (databases.isEmpty())
		{
			this.databaseMap = Collections.emptySortedMap();
		}
		else if (databases.size() == 1)
		{
			this.databaseMap = Collections.singletonSortedMap(databases.iterator().next(), new AtomicInteger(1));
		}
		else
		{
			SortedMap map = new TreeMap();
			
			for (D database: databases)
			{
				map.put(database, new AtomicInteger(1));
			}
			
			this.databaseMap = map;
		}
	}
	
	/**
	 * {@inheritDoc}
	 * @see net.sf.hajdbc.balancer.Balancer#primary()
	 */
	@Override
	public D primary()
	{
		try
		{
			return this.databaseMap.firstKey();
		}
		catch (NoSuchElementException e)
		{
			return null;
		}
	}

	@Override
	public D local() {
		try
		{
			return getDatabases().stream()
					.filter(e->e.isLocal()).findFirst()
					.orElse(null);
		}
		catch (NoSuchElementException e)
		{
			return null;
		}
	}

	/**
	 * {@inheritDoc}
	 * @see net.sf.hajdbc.balancer.AbstractBalancer#getDatabases()
	 */
	@Override
	public Set getDatabases()
	{
		return this.databaseMap.keySet();
	}

	/**
	 * {@inheritDoc}
	 * @see java.util.Set#addAll(java.util.Collection)
	 */
	@Override
	public boolean addAll(Collection databases)
	{
		this.lock.lock();
		
		try
		{
			SortedMap addMap = new TreeMap(this.databaseMap);
			
			boolean added = false;
			
			for (D database: databases)
			{
				added = (addMap.put(database, new AtomicInteger(1)) == null) || added;
			}
			
			if (added)
			{
				this.databaseMap = addMap;
			}
			
			return added;
		}
		finally
		{
			this.lock.unlock();
		}
	}

	/**
	 * {@inheritDoc}
	 * @see java.util.Set#removeAll(java.util.Collection)
	 */
	@Override
	public boolean removeAll(Collection databases)
	{
		this.lock.lock();
		
		try
		{
			SortedMap map = new TreeMap(this.databaseMap);
			
			boolean removed = map.keySet().removeAll(databases);

			if (removed)
			{
				this.databaseMap = map;
			}
			
			return removed;
		}
		finally
		{
			this.lock.unlock();
		}
	}

	/**
	 * {@inheritDoc}
	 * @see java.util.Set#retainAll(java.util.Collection)
	 */
	@Override
	public boolean retainAll(Collection databases)
	{
		this.lock.lock();
		
		try
		{
			SortedMap map = new TreeMap(this.databaseMap);
			
			boolean retained = map.keySet().retainAll(databases);

			if (retained)
			{
				this.databaseMap = map;
			}
			
			return retained;
		}
		finally
		{
			this.lock.unlock();
		}
	}

	/**
	 * {@inheritDoc}
	 * @see java.util.Set#clear()
	 */
	@Override
	public void clear()
	{
		this.lock.lock();
		
		try
		{
			if (!this.databaseMap.isEmpty())
			{
				this.databaseMap = new TreeMap();
			}
		}
		finally
		{
			this.lock.unlock();
		}
	}

	/**
	 * {@inheritDoc}
	 * @see java.util.Set#remove(java.lang.Object)
	 */
	@Override
	public boolean remove(Object database)
	{
		this.lock.lock();
		
		try
		{
			boolean remove = this.databaseMap.containsKey(database);
			
			if (remove)
			{
				if (this.databaseMap.size() == 1)
				{
					this.databaseMap = Collections.emptySortedMap();
				}
				else
				{
					SortedMap map = new TreeMap(this.databaseMap);

					map.remove(database);
					
					this.databaseMap = map;
				}
			}
			
			return remove;
		}
		finally
		{
			this.lock.unlock();
		}
	}

	/**
	 * {@inheritDoc}
	 * @see net.sf.hajdbc.balancer.Balancer#next()
	 */
	@Override
	public D next()
	{
		Set> entrySet = this.databaseMap.entrySet();
		
		return !entrySet.isEmpty() ? java.util.Collections.min(entrySet, this.comparator).getKey() : null;
	}

	/**
	 * {@inheritDoc}
	 * @see java.util.Set#add(java.lang.Object)
	 */
	@Override
	public boolean add(D database)
	{
		this.lock.lock();
		
		try
		{
			boolean add = !this.databaseMap.containsKey(database);
			
			if (add)
			{
				AtomicInteger load = new AtomicInteger(1);
				
				if (this.databaseMap.isEmpty())
				{
					this.databaseMap = Collections.singletonSortedMap(database, load);
				}
				else
				{
					SortedMap map = new TreeMap(this.databaseMap);
					
					map.put(database, load);
					
					this.databaseMap = map;
				}
			}
			
			return add;
		}
		finally
		{
			this.lock.unlock();
		}
	}
	
	/**
	 * {@inheritDoc}
	 * @see net.sf.hajdbc.balancer.Balancer#invoke(net.sf.hajdbc.invocation.Invoker, net.sf.hajdbc.Database, java.lang.Object)
	 */
	@Override
	public  R invoke(Invoker invoker, D database, T object) throws E
	{
		AtomicInteger load = this.databaseMap.get(database);
		
		if (load != null)
		{
			load.incrementAndGet();
		}
		
		try
		{
			return invoker.invoke(database, object);
		}
		finally
		{
			if (load != null)
			{
				load.decrementAndGet();
			}
		}
	}
}