001/*
002 * HA-JDBC: High-Availability JDBC
003 * Copyright (c) 2004-2007 Paul Ferraro
004 * 
005 * This library is free software; you can redistribute it and/or modify it 
006 * under the terms of the GNU Lesser General Public License as published by the 
007 * Free Software Foundation; either version 2.1 of the License, or (at your 
008 * option) any later version.
009 * 
010 * This library is distributed in the hope that it will be useful, but WITHOUT
011 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or 
012 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU Lesser General Public License 
013 * for more details.
014 * 
015 * You should have received a copy of the GNU Lesser General Public License
016 * along with this library; if not, write to the Free Software Foundation, 
017 * Inc., 59 Temple Place, Suite 330, Boston, MA 02111-1307 USA
018 * 
019 * Contact: ferraro@users.sourceforge.net
020 */
021package net.sf.hajdbc.balancer;
022
023import java.util.Collections;
024import java.util.Comparator;
025import java.util.Map;
026import java.util.Set;
027import java.util.TreeMap;
028import java.util.concurrent.atomic.AtomicInteger;
029import java.util.concurrent.locks.Lock;
030import java.util.concurrent.locks.ReentrantLock;
031
032import net.sf.hajdbc.Balancer;
033import net.sf.hajdbc.Database;
034
035/**
036 * Balancer implementation whose {@link #next()} implementation returns the database with the least load.
037 *
038 * @author  Paul Ferraro
039 * @param <D> either java.sql.Driver or javax.sql.DataSource
040 */
041public class LoadBalancer<D> implements Balancer<D>
042{
043        private volatile Map<Database<D>, AtomicInteger> databaseMap = Collections.emptyMap();
044
045        private Lock lock = new ReentrantLock();
046        
047        private Comparator<Map.Entry<Database<D>, AtomicInteger>> comparator = new Comparator<Map.Entry<Database<D>, AtomicInteger>>()
048        {
049                public int compare(Map.Entry<Database<D>, AtomicInteger> mapEntry1, Map.Entry<Database<D>, AtomicInteger> mapEntry2)
050                {
051                        Database<D> database1 = mapEntry1.getKey();
052                        Database<D> database2 = mapEntry2.getKey();
053
054                        float load1 = mapEntry1.getValue().get();
055                        float load2 = mapEntry2.getValue().get();
056                        
057                        int weight1 = database1.getWeight();
058                        int weight2 = database2.getWeight();
059                        
060                        // If weights are the same, we can simply compare the loads
061                        if (weight1 == weight2)
062                        {
063                                return Float.compare(load1, load2);
064                        }
065                        
066                        float weightedLoad1 = (weight1 != 0) ? (load1 / weight1) : Float.POSITIVE_INFINITY;
067                        float weightedLoad2 = (weight2 != 0) ? (load2 / weight2) : Float.POSITIVE_INFINITY;
068                        
069                        return Float.compare(weightedLoad1, weightedLoad2);
070                }
071        };
072        
073        /**
074         * @see net.sf.hajdbc.Balancer#all()
075         */
076        @Override
077        public Set<Database<D>> all()
078        {
079                return Collections.unmodifiableSet(this.databaseMap.keySet());
080        }
081
082        /**
083         * @see net.sf.hajdbc.Balancer#clear()
084         */
085        @Override
086        public void clear()
087        {
088                this.lock.lock();
089                
090                try
091                {
092                        this.databaseMap = Collections.emptyMap();
093                }
094                finally
095                {
096                        this.lock.unlock();
097                }
098        }
099
100        /**
101         * @see net.sf.hajdbc.Balancer#remove(net.sf.hajdbc.Database)
102         */
103        @Override
104        public boolean remove(Database<D> database)
105        {
106                this.lock.lock();
107                
108                try
109                {
110                        boolean exists = this.databaseMap.containsKey(database);
111                        
112                        if (exists)
113                        {
114                                Map<Database<D>, AtomicInteger> map = new TreeMap<Database<D>, AtomicInteger>(this.databaseMap);
115                                
116                                map.remove(database);
117
118                                this.databaseMap = map;
119                        }
120                        
121                        return exists;
122                }
123                finally
124                {
125                        this.lock.unlock();
126                }
127        }
128
129        /**
130         * @see net.sf.hajdbc.Balancer#next()
131         */
132        @Override
133        public Database<D> next()
134        {
135                return Collections.min(this.databaseMap.entrySet(), this.comparator).getKey();
136        }
137
138        /**
139         * @see net.sf.hajdbc.Balancer#add(net.sf.hajdbc.Database)
140         */
141        @Override
142        public boolean add(Database<D> database)
143        {
144                this.lock.lock();
145                
146                try
147                {
148                        boolean exists = this.databaseMap.containsKey(database);
149                        
150                        if (!exists)
151                        {
152                                Map<Database<D>, AtomicInteger> map = new TreeMap<Database<D>, AtomicInteger>(this.databaseMap);
153                                
154                                map.put(database, new AtomicInteger(1));
155
156                                this.databaseMap = map;
157                        }
158                        
159                        return !exists;
160                }
161                finally
162                {
163                        this.lock.unlock();
164                }
165        }
166        
167        /**
168         * @see net.sf.hajdbc.Balancer#beforeInvocation(net.sf.hajdbc.Database)
169         */
170        @Override
171        public void beforeInvocation(Database<D> database)
172        {
173                AtomicInteger load = this.databaseMap.get(database);
174                
175                if (load != null)
176                {
177                        load.incrementAndGet();
178                }
179        }
180        
181        /**
182         * @see net.sf.hajdbc.Balancer#afterInvocation(net.sf.hajdbc.Database)
183         */
184        @Override
185        public void afterInvocation(Database<D> database)
186        {
187                AtomicInteger load = this.databaseMap.get(database);
188                
189                if (load != null)
190                {
191                        load.decrementAndGet();
192                }
193        }
194}