Cache.java

/*
 * Copyright 2012 LinkedIn Corp.
 * 
 * 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 azkaban.utils.cache;

import java.util.ArrayList;
import java.util.Map;
import java.util.concurrent.ConcurrentHashMap;

public class Cache {
	private long nextUpdateTime = 0;
	private long updateFrequency = 1 * 60 * 1000;
	private int maxCacheSize = -1;

	private long expireTimeToLive = -1; // Never expires
	private long expireTimeToIdle = -1;

	private EjectionPolicy ejectionPolicy = EjectionPolicy.LRU;
	private CacheManager manager = null;

	private Map<Object, Element<?>> elementMap = new ConcurrentHashMap<Object, Element<?>>();

	public enum EjectionPolicy {
		LRU, FIFO
	}

	/* package */Cache(CacheManager manager) {
		this.manager = manager;
	}

	@SuppressWarnings("unchecked")
	public <T> T get(Object key) {
		Element<?> element = elementMap.get(key);
		if (element == null) {
			return null;
		}
		return (T) element.getElement();
	}

	public <T> void put(Object key, T item) {
		Element<T> elem = new Element<T>(key, item);
		elementMap.put(key, elem);
	}

	public boolean remove(Object key) {
		Element<?> elem = elementMap.remove(key);
		if (elem == null) {
			return false;
		}
		
		return true;
	}

	public Cache setMaxCacheSize(int size) {
		maxCacheSize = size;
		return this;
	}

	public Cache setEjectionPolicy(EjectionPolicy policy) {
		ejectionPolicy = policy;
		return this;
	}

	public Cache setUpdateFrequencyMs(long updateFrequencyMs) {
		this.updateFrequency = updateFrequencyMs;
		return this;
	}

	public Cache setExpiryTimeToLiveMs(long time) {
		this.expireTimeToLive = time;
		if (time > 0) {
			manager.update();
		}

		return this;
	}

	public Cache setExpiryIdleTimeMs(long time) {
		this.expireTimeToIdle = time;
		if (time > 0) {
			manager.update();
		}
		return this;
	}

	public int getSize() {
		return elementMap.size();
	}

	public long getExpireTimeToLive() {
		return expireTimeToLive;
	}

	public long getExpireTimeToIdle() {
		return expireTimeToIdle;
	}

	public synchronized <T> void insertElement(Object key, T item) {
		if (maxCacheSize < 0 || elementMap.size() < maxCacheSize) {
			Element<T> elem = new Element<T>(key, item);
			elementMap.put(key, elem);
		} else {
			internalExpireCache();

			Element<T> elem = new Element<T>(key, item);
			if (elementMap.size() < maxCacheSize) {
				elementMap.put(key, elem);
			} else {
				Element<?> element = getNextExpiryElement();
				if (element != null) {
					elementMap.remove(element.getKey());
				}

				elementMap.put(key, elem);
			}
		}
	}

	private Element<?> getNextExpiryElement() {
		if (ejectionPolicy == EjectionPolicy.LRU) {
			long latestAccessTime = Long.MAX_VALUE;
			Element<?> ejectionCandidate = null;
			for (Element<?> elem : elementMap.values()) {
				if (latestAccessTime > elem.getLastUpdateTime()) {
					latestAccessTime = elem.getLastUpdateTime();
					ejectionCandidate = elem;
				}
			}

			return ejectionCandidate;
		} else if (ejectionPolicy == EjectionPolicy.FIFO) {
			long earliestCreateTime = Long.MAX_VALUE;
			Element<?> ejectionCandidate = null;
			for (Element<?> elem : elementMap.values()) {
				if (earliestCreateTime > elem.getCreationTime()) {
					earliestCreateTime = elem.getCreationTime();
					ejectionCandidate = elem;
				}
			}
			return ejectionCandidate;
		}

		return null;
	}

	public synchronized void expireCache() {
		long currentTime = System.currentTimeMillis();
		if (nextUpdateTime < currentTime) {
			internalExpireCache();
			nextUpdateTime = currentTime + updateFrequency;
		}
	}

	private synchronized void internalExpireCache() {
		ArrayList<Element<?>> elems = new ArrayList<Element<?>>(elementMap.values());
		
		for (Element<?> elem : elems) {
			if (shouldExpire(elem)) {
				elementMap.remove(elem.getKey());
			}
		}
	}

	private boolean shouldExpire(Element<?> elem) {
		if (expireTimeToLive > -1) {
			if (System.currentTimeMillis() - elem.getCreationTime() > expireTimeToLive) {
				return true;
			}
		}
		if (expireTimeToIdle > -1) {
			if (System.currentTimeMillis() - elem.getLastUpdateTime() > expireTimeToIdle) {
				return true;
			}
		}

		return false;
	}
}