Source code

001/*
002 * HA-JDBC: High-Availability JDBC
003 * Copyright (C) 2012  Paul Ferraro
004 *
005 * This program is free software: you can redistribute it and/or modify
006 * it under the terms of the GNU Lesser General Public License as published by
007 * the Free Software Foundation, either version 3 of the License, or
008 * (at your option) any later version.
009 *
010 * This program is distributed in the hope that it will be useful,
011 * but WITHOUT ANY WARRANTY; without even the implied warranty of
012 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
013 * GNU Lesser General Public License for more details.
014 *
015 * You should have received a copy of the GNU Lesser General Public License
016 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
017 */
018package net.sf.hajdbc.balancer.load;
019
020import java.util.Collection;
021import java.util.Comparator;
022import java.util.Map;
023import java.util.NoSuchElementException;
024import java.util.Set;
025import java.util.SortedMap;
026import java.util.TreeMap;
027import java.util.concurrent.atomic.AtomicInteger;
028import java.util.concurrent.locks.Lock;
029import java.util.concurrent.locks.ReentrantLock;
030
031import net.sf.hajdbc.Database;
032import net.sf.hajdbc.balancer.AbstractBalancer;
033import net.sf.hajdbc.invocation.Invoker;
034import net.sf.hajdbc.util.Collections;
035
036/**
037 * Balancer implementation whose {@link #next()} implementation returns the database with the least load.
038 *
039 * @author  Paul Ferraro
040 * @param <D> either java.sql.Driver or javax.sql.DataSource
041 */
042public class LoadBalancer<Z, D extends Database<Z>> extends AbstractBalancer<Z, D>
043{
044        private final Lock lock = new ReentrantLock();
045        
046        private volatile SortedMap<D, AtomicInteger> databaseMap = Collections.emptySortedMap();
047        
048        private Comparator<Map.Entry<D, AtomicInteger>> comparator = new Comparator<Map.Entry<D, AtomicInteger>>()
049        {
050                @Override
051                public int compare(Map.Entry<D, AtomicInteger> mapEntry1, Map.Entry<D, AtomicInteger> mapEntry2)
052                {
053                        D database1 = mapEntry1.getKey();
054                        D database2 = mapEntry2.getKey();
055
056                        float load1 = mapEntry1.getValue().get();
057                        float load2 = mapEntry2.getValue().get();
058                        
059                        int weight1 = database1.getWeight();
060                        int weight2 = database2.getWeight();
061                        
062                        // If weights are the same, we can simply compare the loads
063                        if (weight1 == weight2)
064                        {
065                                return Float.compare(load1, load2);
066                        }
067                        
068                        float weightedLoad1 = (weight1 != 0) ? (load1 / weight1) : Float.POSITIVE_INFINITY;
069                        float weightedLoad2 = (weight2 != 0) ? (load2 / weight2) : Float.POSITIVE_INFINITY;
070                        
071                        return Float.compare(weightedLoad1, weightedLoad2);
072                }
073        };
074
075        /**
076         * Constructs a new LoadBalancer
077         * @param databases
078         */
079        public LoadBalancer(Set<D> databases)
080        {
081                if (databases.isEmpty())
082                {
083                        this.databaseMap = Collections.emptySortedMap();
084                }
085                else if (databases.size() == 1)
086                {
087                        this.databaseMap = Collections.singletonSortedMap(databases.iterator().next(), new AtomicInteger(1));
088                }
089                else
090                {
091                        SortedMap<D, AtomicInteger> map = new TreeMap<D, AtomicInteger>();
092                        
093                        for (D database: databases)
094                        {
095                                map.put(database, new AtomicInteger(1));
096                        }
097                        
098                        this.databaseMap = map;
099                }
100        }
101        
102        /**
103         * {@inheritDoc}
104         * @see net.sf.hajdbc.balancer.Balancer#primary()
105         */
106        @Override
107        public D primary()
108        {
109                try
110                {
111                        return this.databaseMap.firstKey();
112                }
113                catch (NoSuchElementException e)
114                {
115                        return null;
116                }
117        }
118
119        /**
120         * {@inheritDoc}
121         * @see net.sf.hajdbc.balancer.AbstractBalancer#getDatabases()
122         */
123        @Override
124        protected Set<D> getDatabases()
125        {
126                return this.databaseMap.keySet();
127        }
128
129        /**
130         * {@inheritDoc}
131         * @see java.util.Set#addAll(java.util.Collection)
132         */
133        @Override
134        public boolean addAll(Collection<? extends D> databases)
135        {
136                this.lock.lock();
137                
138                try
139                {
140                        SortedMap<D, AtomicInteger> addMap = new TreeMap<D, AtomicInteger>(this.databaseMap);
141                        
142                        boolean added = false;
143                        
144                        for (D database: databases)
145                        {
146                                added = (addMap.put(database, new AtomicInteger(1)) == null) || added;
147                        }
148                        
149                        if (added)
150                        {
151                                this.databaseMap = addMap;
152                        }
153                        
154                        return added;
155                }
156                finally
157                {
158                        this.lock.unlock();
159                }
160        }
161
162        /**
163         * {@inheritDoc}
164         * @see java.util.Set#removeAll(java.util.Collection)
165         */
166        @Override
167        public boolean removeAll(Collection<?> databases)
168        {
169                this.lock.lock();
170                
171                try
172                {
173                        SortedMap<D, AtomicInteger> map = new TreeMap<D, AtomicInteger>(this.databaseMap);
174                        
175                        boolean removed = map.keySet().removeAll(databases);
176
177                        if (removed)
178                        {
179                                this.databaseMap = map;
180                        }
181                        
182                        return removed;
183                }
184                finally
185                {
186                        this.lock.unlock();
187                }
188        }
189
190        /**
191         * {@inheritDoc}
192         * @see java.util.Set#retainAll(java.util.Collection)
193         */
194        @Override
195        public boolean retainAll(Collection<?> databases)
196        {
197                this.lock.lock();
198                
199                try
200                {
201                        SortedMap<D, AtomicInteger> map = new TreeMap<D, AtomicInteger>(this.databaseMap);
202                        
203                        boolean retained = map.keySet().retainAll(databases);
204
205                        if (retained)
206                        {
207                                this.databaseMap = map;
208                        }
209                        
210                        return retained;
211                }
212                finally
213                {
214                        this.lock.unlock();
215                }
216        }
217
218        /**
219         * {@inheritDoc}
220         * @see java.util.Set#clear()
221         */
222        @Override
223        public void clear()
224        {
225                this.lock.lock();
226                
227                try
228                {
229                        if (!this.databaseMap.isEmpty())
230                        {
231                                this.databaseMap = new TreeMap<D, AtomicInteger>();
232                        }
233                }
234                finally
235                {
236                        this.lock.unlock();
237                }
238        }
239
240        /**
241         * {@inheritDoc}
242         * @see java.util.Set#remove(java.lang.Object)
243         */
244        @Override
245        public boolean remove(Object database)
246        {
247                this.lock.lock();
248                
249                try
250                {
251                        boolean remove = this.databaseMap.containsKey(database);
252                        
253                        if (remove)
254                        {
255                                if (this.databaseMap.size() == 1)
256                                {
257                                        this.databaseMap = Collections.emptySortedMap();
258                                }
259                                else
260                                {
261                                        SortedMap<D, AtomicInteger> map = new TreeMap<D, AtomicInteger>(this.databaseMap);
262
263                                        map.remove(database);
264                                        
265                                        this.databaseMap = map;
266                                }
267                        }
268                        
269                        return remove;
270                }
271                finally
272                {
273                        this.lock.unlock();
274                }
275        }
276
277        /**
278         * {@inheritDoc}
279         * @see net.sf.hajdbc.balancer.Balancer#next()
280         */
281        @Override
282        public D next()
283        {
284                Set<Map.Entry<D, AtomicInteger>> entrySet = this.databaseMap.entrySet();
285                
286                return !entrySet.isEmpty() ? java.util.Collections.min(entrySet, this.comparator).getKey() : null;
287        }
288
289        /**
290         * {@inheritDoc}
291         * @see java.util.Set#add(java.lang.Object)
292         */
293        @Override
294        public boolean add(D database)
295        {
296                this.lock.lock();
297                
298                try
299                {
300                        boolean add = !this.databaseMap.containsKey(database);
301                        
302                        if (add)
303                        {
304                                AtomicInteger load = new AtomicInteger(1);
305                                
306                                if (this.databaseMap.isEmpty())
307                                {
308                                        this.databaseMap = Collections.singletonSortedMap(database, load);
309                                }
310                                else
311                                {
312                                        SortedMap<D, AtomicInteger> map = new TreeMap<D, AtomicInteger>(this.databaseMap);
313                                        
314                                        map.put(database, load);
315                                        
316                                        this.databaseMap = map;
317                                }
318                        }
319                        
320                        return add;
321                }
322                finally
323                {
324                        this.lock.unlock();
325                }
326        }
327        
328        /**
329         * {@inheritDoc}
330         * @see net.sf.hajdbc.balancer.Balancer#invoke(net.sf.hajdbc.invocation.Invoker, net.sf.hajdbc.Database, java.lang.Object)
331         */
332        @Override
333        public <T, R, E extends Exception> R invoke(Invoker<Z, D, T, R, E> invoker, D database, T object) throws E
334        {
335                AtomicInteger load = this.databaseMap.get(database);
336                
337                if (load != null)
338                {
339                        load.incrementAndGet();
340                }
341                
342                try
343                {
344                        return invoker.invoke(database, object);
345                }
346                finally
347                {
348                        if (load != null)
349                        {
350                                load.decrementAndGet();
351                        }
352                }
353        }
354}