PriorityBufferMap.java

/***************************************************************************
   Copyright 2008 Emily Estes

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License 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 net.metanotion.util;


import java.util.Iterator;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.atomic.AtomicLong;

/** This class implements a Map which ejects elements if it gets too "big".
It is assumed that entries implement the {@link net.metanotion.util.Cacheable} interface which allows the entry to
declare it's "weight". Entries are ejected from the map based on a Least-Recently-Used basis until the sum of the weights
is less than the maximum allowed for the map instance.
	@param <K> The type of keys this dictionary uses as an index.
	@param <V> The type of values this dictionary stores.
*/
public final class PriorityBufferMap<K,V extends Cacheable> implements MutableDictionary<K,V>, Iterable<Map.Entry<K,V>> {
	private final Map<K,LinkEl> map = new ConcurrentHashMap<K,LinkEl>();
	private final long maxWeight;

	private final AtomicLong curWeight = new AtomicLong(0);
	protected LinkEl first = null;
	protected LinkEl last = null;

	/** Create a new map with a maximum allowed weight.
		@param maxWeight The maximum weight of cached objects until entries are ejected.
	*/
	public PriorityBufferMap(final long maxWeight) {
		this.maxWeight = maxWeight;
		if(maxWeight < 0) { throw new IllegalArgumentException("maxWeight must be greater than or equal to 0"); }
	}

	@Override public Iterator<Map.Entry<K,V>> iterator() { return new PriorityBufferIterator(first); }

	private void reduce() {
		if(last == null) { return; }
		final LinkEl le = last;
		map.remove(le.key);
		curWeight.addAndGet(-1 * le.val.weight());
		if(le.prev != null) { last.prev.next = null; }
		last = le.prev;
		try {
			le.val.close();
		} catch (RuntimeException re) { }
	}

	private void moveToFront(LinkEl le) {
		if(le != first) {
			final LinkEl p = le.prev;
			le.prev.next = le.next;
			if(le.next != null) {
				le.next.prev = p;
			} else {
				last = p;
			}
			le.prev = null;
			le.next = first;
			first.prev = le;
			first = le;
		}
	}

	@Override public void put(final K k, final V v) {
		LinkEl le = map.get(k);
		V ret = null;
		long cw;
		if(le == null) {
			le = new LinkEl(k,v,null,first);
			if(map.size() == 0) {
				last = le;
			} else {
				first.prev = le;
			}
			first = le;
			map.put(k, le);
			cw = curWeight.addAndGet(v.weight());
		} else {
			ret = le.val;
			cw = curWeight.addAndGet(v.weight() - le.val.weight());
			le.val = v;
			moveToFront(le);
		}
		if(cw > maxWeight) { reduce(); }
		if(ret != null) { ret.close(); }
	}

	@Override public V remove(final K k) {
		final LinkEl le = map.remove(k);
		if(le == null) { return null; }

		curWeight.addAndGet(-1 * le.val.weight());
		if(le.prev != null) { le.prev.next = le.next; }
		if(le.next != null) { le.next.prev = le.prev; }
		if(le.prev == null) { first = le.next; }
		if(le.next == null) { last = le.prev; }

		return le.val;
	}

	@Override public V get(final K k) {
		final LinkEl le = map.get(k);
		if(le == null) { return null; }
		moveToFront(le);
		return le.val;
	}

	public boolean isEmpty() { return map.isEmpty(); }
	public int size() { return map.size(); }

	/** Get the current weight of the map.
		@return The sum of the weights from the cached objects in the map.
	*/
	public long weight()		{ return curWeight.get(); }

	// Utility classes.
	private final class LinkEl implements Map.Entry<K,V> {
		public LinkEl next;
		public LinkEl prev;
		public final K key;
		public V val;

		public LinkEl(final K k, final V v, final LinkEl prev, final LinkEl next) {
			this.key = k;
			this.val = v;
			this.prev = prev;
			this.next = next;
		}

		@Override public boolean equals(final Object o) {
			if(!(o instanceof Map.Entry)) { return false; }
			final Map.Entry me = (Map.Entry) o;
			if(!key.equals(me.getKey())) { return false; }
			if(!val.equals(me.getValue())) { return false; }
			return true;
		}

		@Override public int hashCode() { return key.hashCode() + val.hashCode(); }
		@Override public K getKey() { return key; }
		@Override public V getValue() { return val; }
		@Override public V setValue(V value) { throw new UnsupportedOperationException(); }
	}


	private final class PriorityBufferIterator implements Iterator<Map.Entry<K,V>> {
		private LinkEl current;
		public PriorityBufferIterator(final LinkEl start) { this.current = start; }
		@Override public boolean hasNext() { return (current != null); }
		@Override public void remove() { throw new UnsupportedOperationException(); }

		@Override public Map.Entry<K,V> next() {
			if(current == null) { throw new java.util.NoSuchElementException(); }
			final LinkEl c = current;
			current = current.next;
			return c;
		}
	}
}