/* 
 * ConnectionsPool.java - Classes that gathers the functionality to
 *                        handle connections with the DB. These classes
 *                        were not included in the original version of
 *                        TPC-W, they were included by the Laboratory
 *                        of Distributed Systems, of the Technical University
 *                        of Madrid.
 ************************************************************************
 *
 * This is part of the modified Java TPC-W distribution. The original TPC-W
 * distribution was developed by Harold Cain, Tim Heil, Milo Martin, Eric
 * Weglarz, and Todd Bezenek.  University of Wisconsin - Madison,
 * Computer Sciences Dept. and Dept. of Electrical and Computer Engineering,
 * as a part of Prof. Mikko Lipasti's Fall 1999 ECE 902 course. The modified
 * TPC-W distribution (this one) adds some enhancements that increases its
 * scalability. The modified TPC-W distribution was developed by the Laboratory
 * of Distributed Systems, of the Technical University of Madrid. This file
 * was not included in the original TPC-W distribution, it was developed 
 * for the modified TPC-W version. 
 *
 * Copyright (C) 2012 by Luis Rodero-Merino.
 *
 * This source code is distributed "as is" in the hope that it will be
 * useful.  It comes with no warranty, and no author or distributor
 * accepts any responsibility for the consequences of its use.
 *
 * Everyone is granted permission to copy, modify and redistribute
 * this code under the following conditions:
 *
 * This code is distributed for non-commercial use only.
 * Please contact the maintainer for restrictions applying to 
 * commercial use of these tools.
 *
 * Permission is granted to anyone to make or distribute copies
 * of this code, either as received or modified, in any
 * medium, provided that all copyright notices, permission and
 * nonwarranty notices are preserved, and that the distributor
 * grants the recipient permission for further redistribution as
 * permitted by this document.
 *
 * Permission is granted to distribute this code in compiled
 * or executable form under the same conditions that apply for
 * source code, provided that either:
 *
 * A. it is accompanied by the corresponding machine-readable
 *    source code,
 * B. it is accompanied by a written offer, with no time limit,
 *    to give anyone a machine-readable copy of the corresponding
 *    source code in return for reimbursement of the cost of
 *    distribution.  This written offer must permit verbatim
 *    duplication by anyone, or
 * C. it is distributed by someone who received only the
 *    executable form, and is accompanied by a copy of the
 *    written offer of source code that they received concurrently.
 *
 * In other words, you are welcome to use, share and improve this codes.
 * You are forbidden to forbid anyone else to use, share and improve what
 * you give them.
 *
 ************************************************************************/
import java.net.InetSocketAddress;
import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.SQLException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.HashSet;
import java.util.Iterator;
import java.util.LinkedHashSet;
import java.util.List;
import java.util.Map;
import java.util.Properties;
import java.util.Set;

public class ConnectionsPool {

	protected static final int DEFAULT_INITIAL_POOL_SIZE = 50;
	protected static final int DEFAULT_THREADS_TO_FILL_POOL = 10;
	
	// Must make sure that iterator traverses the set in the same order connections were added, this is required by method getConnection()
	protected Set<Connection> connectionsAvailable = new LinkedHashSet<Connection>(); 
	protected Set<Connection> connectionsInUse = new HashSet<Connection>();
	
	// Whenever the set of database replicas known (i.e. the 'view') changes we must
	// redo the connections in order to benefit from new replicas or reconnect when some
	// replica is gone. These two list store:
	//  1- The connections that have to be closed (they are connected to some replica that is not there anymore)
	protected Set<Connection> connectionsInUseToRemove = new HashSet<Connection>();
	//  2- The connections that have to be redone 
	protected Set<Connection> connectionsInUseToRedo = new HashSet<Connection>();
	
	protected String jdbcPath = null; 
	protected String jdbcDriverClass = null;
	protected int connsAmount = -1;
	protected int threadsAmount = -1;
	protected Properties driverConf = null;

	protected boolean useSessions = false;
	// Mapping so the same session always gets connected to the same host, used to implement sticky sessions between servlet and database,
	// ignored if no pool is used
	private Map<String, String> sessionsHostsMap = new HashMap<String, String>(); 
	
	protected static ConnectionsPool singletonInstance = null;
	
	public static synchronized void preparePool(String jdbcPath, String jdbcDriverClass, boolean useSessions) throws SQLException {
		preparePool(jdbcPath, jdbcDriverClass, useSessions, null);
	}
	
	public static synchronized void preparePool(String jdbcPath, String jdbcDriverClass, boolean useSessions, Properties driverConf) throws SQLException {
		preparePool(jdbcPath, jdbcDriverClass, useSessions, driverConf, DEFAULT_INITIAL_POOL_SIZE, DEFAULT_THREADS_TO_FILL_POOL);
	}
	
	public static synchronized void preparePool(String jdbcPath, String jdbcDriverClass, boolean useSessions, Properties driverConf, int connsAmount, int threadsAmount) throws SQLException {
		if(singletonInstance != null) {
			System.err.println("Ignoring petition to prepare pool, it was already prepared!");
			return;
		}
		singletonInstance = new ConnectionsPool(jdbcPath, jdbcDriverClass, useSessions, driverConf, connsAmount, threadsAmount);
	}
	
	protected ConnectionsPool(String jdbcPath, String jdbcDriverClass, boolean useSessions, Properties driverConf, int connsAmount, int threadsAmount) throws SQLException {
		this.jdbcPath = jdbcPath;
		this.jdbcDriverClass = jdbcDriverClass;
		this.useSessions = useSessions;
		this.driverConf = driverConf;
		this.connsAmount = connsAmount;
		this.threadsAmount = threadsAmount;
		fillPool();
	}
	
	public static ConnectionsPool getInstance() {
		if(singletonInstance == null)
			throw new IllegalStateException("Must prepare pool before trying to use it");
		return singletonInstance;
	}
	
	private synchronized void fillPool() throws SQLException {		

		System.out.println("Loading JDBC driver '" + jdbcDriverClass + "'");
		
		try {
			Class.forName(jdbcDriverClass);
		} catch (ClassNotFoundException exception) {
			String errMsg = "It seems the postgresql driver (class name ' " + jdbcDriverClass + "') cannot be found, check the proper .jar file is available";
			System.err.println(errMsg);
			throw new SQLException(errMsg);
		}

		System.out.println("JDBC driver '" + jdbcDriverClass + "' loaded");
		System.out.println("Filling connections pool, creating " + connsAmount + " connections with " + threadsAmount + " threads");
		connectionsAvailable = new PoolFiller(this, connsAmount, threadsAmount).fill();
	}

	public synchronized void returnConnection(Connection conn) throws SQLException {
		
		if(conn == null)
			return;

		if(connectionsAvailable.contains(conn))
			throw new IllegalStateException("A connection that was supposed to be available was returned!");
		
		// We only put back the connection in the available connections list if it was created by us
		if(!connectionsInUse.remove(conn))
			return;
		
		if(conn.isClosed())
			return;
		
		if(connectionsInUseToRemove.contains(conn)) {
			conn.close();
			return;
		}
		
		connectionsAvailable.add(conn);
		
	}
    
    // Get a connection from the pool.
    public synchronized Connection getConnection() throws SQLException {
    	
	    Connection conn = null;
	    
	    while (!connectionsAvailable.isEmpty()) {
	    	// Pick the first Connection in the list
			// to get round-robin usage
	    	Iterator<Connection> connsAvIter = connectionsAvailable.iterator();
	    	conn = (Connection)connsAvIter.next();
	    	connsAvIter.remove();
	    	if (conn.isClosed())
    			continue;
			// Got a connection.
	    	connectionsInUse.add(conn); 
	    	return conn;
	    }

		conn = createConnection();	
	    if(conn != null)
	    	connectionsInUse.add(conn);
	    
	    return conn;
    }
    
    // Get a connection from the pool, associated to some session (identified by its session id)
    public synchronized Connection getConnection(String sessionId) throws SQLException {

    	// If session has already one host associated, then we must look for connections to that host 
		// among available connections
		String host = sessionsHostsMap.get(sessionId);
		
		if(host == null) {
			Connection conn = getConnection();
		    if(conn != null)
				sessionsHostsMap.put(sessionId, conn.getMetaData().getURL());
			return conn;
		}

		List<Connection> connectionsToHost = new ArrayList<Connection>(connectionsAvailable);
		
		Iterator<Connection> connsIter = connectionsToHost.iterator();
		while(connsIter.hasNext())
			if(!(connsIter.next().getMetaData().getURL().equals(host)))
				connsIter.remove();
		
		while(!connectionsToHost.isEmpty()) {
	    	// Pick the first Connection in the list
			// to get round-robin usage
	    	Connection conn = connectionsToHost.remove(0);
	    	connectionsAvailable.remove(conn);
    		if (conn.isClosed())
    			continue;
	    	connectionsInUse.add(conn);
	    	return conn;
		}
	    
		// Problem, no connection available to host... we must create one
		Connection conn = createConnection( host.startsWith("jdbc:") ? host : "jdbc:"+host);	
	    if(conn != null) {
	    	connectionsInUse.add(conn);
	    	sessionsHostsMap.put(sessionId, conn.getMetaData().getURL());
	    }

	    return conn;
    }
    
    protected Connection createConnection() throws SQLException {
	    return createConnection(jdbcPath);
    }
    
    protected Connection createConnection(String path) throws SQLException {
    	
	    Connection conn = null;
    	if(driverConf != null)
    		conn = DriverManager.getConnection(path, driverConf);
    	else
    		conn = DriverManager.getConnection(path);
    	
    	if(conn != null)
    		conn.setAutoCommit(false);
    	
    	if(conn != null) {
    		System.out.println("[" + Thread.currentThread().getName() + "] ConnectionsPool > CREATED CONNECTION TO " + conn.getMetaData().getURL());
    	} else {
    		System.out.println("[" + Thread.currentThread().getName() + "] ConnectionsPool > COULD NOT CREATE CONNECTION");
    	}
    	
	    return conn;    	
    }
	
	protected static InetSocketAddress extractAddress(Connection conn) {

		String connURL = null;
		try {
			connURL = new String(conn.getMetaData().getURL());
		} catch (SQLException exception) {
			System.out.println("[" + Thread.currentThread().getName() + "] ConnectionsPool > ERROR, could not retrieve URL from present SQL connection (?)");
    		return null;	
		}
		// The URL has the following structure 'protocol://host:port/dbpath'
		// Unfortunately, it seems that neither the java.net.URI nor the java.net.URL classes can parse
		// jdbc URLs. So we have to do it manually.
		// Removing the protocol part at the beginning of the URL
		String protocolDelimiter = "://";
		if(connURL.indexOf(protocolDelimiter) >= 0)
			connURL = connURL.substring(connURL.indexOf(protocolDelimiter)+protocolDelimiter.length());
		// Removing the dbpath at the end of the URL
		String dbpathDelimiter = "/";
		if(connURL.indexOf(dbpathDelimiter) >= 0)
			connURL = connURL.substring(0, connURL.indexOf(dbpathDelimiter));
		String[] hostAndPort = connURL.split(":");
		String host = hostAndPort[0];
		int port = -1;
		try {
			port = Integer.parseInt(hostAndPort[1]);
		} catch(NumberFormatException exception) {
			System.out.println("[" + Thread.currentThread().getName() + "] ConnectionsPool > ERROR, could not parse database URL 'conn.getMetaData().getURL()'");
			return null;
		}				
		
		return new InetSocketAddress(host, port);
	}
	
}

class PoolFiller {
	
	private ConnectionsPool connectionsPool = null;
	private int connsAmount = 0;
	private int threadsAmount = 0;
	private int threadsRunning = 0;
	private Set<Connection> connections = new HashSet<Connection>();
	
	protected PoolFiller(ConnectionsPool connectionsPool, int connsAmount, int threadsAmount) {
		
		if((connsAmount <= 0) || (threadsAmount <= 0))
			return;
	
		if(connsAmount < threadsAmount)
			threadsAmount = connsAmount;
		
		this.connsAmount = connsAmount;
		this.threadsAmount = threadsAmount;
		this.connectionsPool = connectionsPool;
	}
	
	protected Set<Connection> fill() {

		int connsPerThread = connsAmount/threadsAmount;
		
		threadsRunning = threadsAmount;
		for(int threadCount = 0; threadCount < threadsAmount-1; threadCount++) 
			new Thread(new ConnectionsCreator(connsPerThread, connectionsPool, this)).start();
		new Thread(new ConnectionsCreator(connsPerThread+(connsAmount%threadsAmount), connectionsPool, this)).start();	
		waitUntilAllThreadsFinish();
		return connections;
	}
	
	protected synchronized void threadFinished(List<Connection> connectionsCreated) {
		connections.addAll(connectionsCreated);
		threadsRunning--;
		notify();
	}
	
	private synchronized void waitUntilAllThreadsFinish() {
		while(threadsRunning > 0)
			try {
				wait();
			} catch (InterruptedException e) {
			}
	}
}

class ConnectionsCreator implements Runnable {

	private ConnectionsPool connectionsPool = null;
	private int connsToCreate = 0;
	private PoolFiller poolFiller = null;
	private List<Connection> connectionsCreated = new ArrayList<Connection>();
	
	public ConnectionsCreator(int connsToCreate, ConnectionsPool connectionsPool, PoolFiller poolFiller) {
		this.connsToCreate = connsToCreate;
		this.connectionsPool = connectionsPool;
		this.poolFiller = poolFiller;
	}
	
	@Override
	public void run() {
		
		System.out.println("[Filler:"+Thread.currentThread().getName()+"] Creating " + connsToCreate + " connections to database");
		
		int connsCount = 0;
		while(connsCount < connsToCreate) {
			Connection conn = null;
			try {
				conn = connectionsPool.createConnection();
				try {
					Thread.sleep(1000);
				} catch (InterruptedException exception) {}
			} catch (SQLException exception) {
				System.err.println("[Filler:"+Thread.currentThread().getName()+"] SQLException caught when filling connection pool, " + connectionsCreated.size() + " conns created so far");
				exception.printStackTrace(System.err);
				break;
			}
			if(conn != null) {
				System.out.println("[Filler:"+Thread.currentThread().getName()+"] New connection ready");
				connectionsCreated.add(conn);
				connsCount++;
			}
		}
		poolFiller.threadFinished(connectionsCreated);
	}
	
}