LockManager.java

package simpledb;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.LinkedList;
import java.util.List;
import java.util.Map;
import java.util.Queue;
import java.util.Set;

public class LockManager {

	private Map<TransactionId, Set<PageId>> tidToPids;
	private Map<PageId, Set<TransactionId>> pidToTids;
	private Map<PageId, Permissions> pagePermissions;

	private Map<Integer, TransactionId> tableIdToTid;
	private Map<TransactionId, Set<TransactionId>> dependencyGraph;
	private Map<Integer, Map<TransactionId, Set<TransactionId>>> fileDependencyGraph;

	public LockManager() {
		tidToPids = new HashMap<>();
		pidToTids = new HashMap<>();
		pagePermissions = new HashMap<>();

		tableIdToTid = new HashMap<>();
		dependencyGraph = new HashMap<>();
		fileDependencyGraph = new HashMap<>();
	}

	public synchronized boolean lock(TransactionId tid, PageId pid, Permissions perm)
			throws TransactionAbortedException {
		if (this.hasLock(tid, pid, perm))
			return true;
		if (!(pagePermissions.containsKey(pid) && pagePermissions.get(pid) == Permissions.READ_WRITE)) {
			if (perm == Permissions.READ_WRITE) {
				if (!pidToTids.containsKey(pid)) {
					if (!tidToPids.containsKey(tid))
						tidToPids.put(tid, new HashSet<>());
					if (!pidToTids.containsKey(pid))
						pidToTids.put(pid, new HashSet<>());
					tidToPids.get(tid).add(pid);
					pidToTids.get(pid).add(tid);
					pagePermissions.put(pid, Permissions.READ_WRITE);
					return true;
				} else if (pidToTids.get(pid).size() == 1 && pidToTids.get(pid).toArray()[0].equals(tid)) {
					pagePermissions.put(pid, Permissions.READ_WRITE);
					return true;
				}
			} else if (perm == Permissions.READ_ONLY) {
				if (!tidToPids.containsKey(tid))
					tidToPids.put(tid, new HashSet<>());
				if (!pidToTids.containsKey(pid))
					pidToTids.put(pid, new HashSet<>());
				tidToPids.get(tid).add(pid);
				pidToTids.get(pid).add(tid);
				pagePermissions.put(pid, Permissions.READ_ONLY);
				return true;
			}
		}
		// add dependency to the graph
		if (!this.dependencyGraph.containsKey(tid))
			this.dependencyGraph.put(tid, new HashSet<>());
		this.dependencyGraph.get(tid).addAll(this.pidToTids.get(pid));
		this.dependencyGraph.get(tid).remove(tid);
		
		if (this.detectDeadlocks(tid)) {
			this.dependencyGraph.remove(tid);
			throw new TransactionAbortedException();
		}
		
		return false;
	}

	public synchronized void unlock(TransactionId tid, PageId pid) {
		if (!hasLock(tid, pid, Permissions.READ_ONLY))
			return;
		
		// update the data structure
		this.tidToPids.get(tid).remove(pid);
		this.pidToTids.get(pid).remove(tid);

		// clean the data structure
		if (tidToPids.get(tid).isEmpty())
			tidToPids.remove(tid);
		if (pidToTids.get(pid).isEmpty()) {
			pidToTids.remove(pid);
			pagePermissions.remove(pid);
		}

		// update dependency graph
		Set<TransactionId> notWaitingTids = new HashSet<>();
		for (TransactionId tidWaiting : this.dependencyGraph.keySet()) {
			this.dependencyGraph.get(tidWaiting).remove(tid);
			if (this.dependencyGraph.get(tidWaiting).isEmpty())
				notWaitingTids.add(tidWaiting);
		}

		// clean the dependency graph
		for (TransactionId notWaitingTid : notWaitingTids) {
			this.dependencyGraph.remove(notWaitingTid);
		}
	}

	public synchronized void unlock(TransactionId tid) {
		if (!tidToPids.containsKey(tid))
			return;

		List<PageId> unlockPids = new ArrayList<>();
		for (PageId pid : tidToPids.get(tid))
			unlockPids.add(pid);
		for (int i = 0; i < unlockPids.size(); i++)
			this.unlock(tid, unlockPids.get(i));
	}

	public synchronized boolean hasLock(TransactionId tid, PageId pid, Permissions perm) {
		return this.tidToPids.containsKey(tid) && tidToPids.get(tid).contains(pid) 
				&& this.pidToTids.containsKey(pid) && pidToTids.get(pid).contains(tid)
				&& this.pagePermissions.containsKey(pid)
				&& (this.pagePermissions.get(pid) == Permissions.READ_WRITE || perm == Permissions.READ_ONLY);
	}

	public synchronized boolean lockFile(TransactionId tid, int tableid) {
		if (!tableIdToTid.containsKey(tableid))
			tableIdToTid.put(tableid, tid);

		return tableIdToTid.get(tableid).equals(tid);
	}

	public synchronized void unlockFile(TransactionId tid, int tableid) {
		if (hasFileLock(tid, tableid)) {
			tableIdToTid.remove(tid);
		}
	}

	public synchronized boolean hasFileLock(TransactionId tid, int tableid) {
		return tableIdToTid.containsKey(tableid) && tableIdToTid.get(tableid).equals(tid);
	}

	private synchronized boolean detectDeadlocks(TransactionId tid) {
		Set<TransactionId> visitedNodes = new HashSet<>();
		Queue<TransactionId> queue = new LinkedList<>();
		queue.add(tid);
		while (!queue.isEmpty()) {
			TransactionId node = queue.remove();
			visitedNodes.add(node);
			if (!this.dependencyGraph.containsKey(node))
				continue;
			for (TransactionId nextNode : this.dependencyGraph.get(node)) {
				if (!visitedNodes.contains(nextNode)) {
					queue.add(nextNode);
				} else if (nextNode.equals(tid)) {
					return true;
				}
			}

		}
		return false;
	}
}