Design a thread-safe connection pool II (Java)

780

Problem description: design a thread-safe Connection Pool implementing the following interfaces:

interface Connection {
    /**
     * Closes the connection or returns it back to the connection pool.
     */
    void close();

    /**
     * Executes the query and returns the result.
     * @param query query to execute.
     * @return result of query execution.
     */
    Object execute (Object query);
}

interface ConnectionPool {
    /**
     * Gets a connection from the pool (if it is available) or creates a new one.
     * @return connection
     */
    Connection getConnection();
}

Number of connections to the database are not limited

Solution:
ConnectionPool interface does not have a method to release a connection, let us create another interface with this method:

public interface ReleaseConnection {
    /**
     * Return SQL connection back to the pool.
     * @param connection SQL connection to return
     */
    void close(java.sql.Connection connection);
}

Let us start from implementation of Connection interface. It needs to have references to SQL connection (to execute queries) and a pool instance (to be able to return SQL connection back to the pool):

import java.sql.PreparedStatement;
import java.sql.SQLException;

public class ConnectionImpl implements Connection {
    private final ConnectionPool pool;
    private final java.sql.Connection connection;

    public ConnectionImpl(ConnectionPool pool, java.sql.Connection connection) {
        this.pool = pool;
        this.connection = connection;
    }

    @Override
    public void close() {
        if (pool instanceof ReleaseConnection) {
            ((ReleaseConnection) pool).close(connection);
        } else {
            try {
                connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public Object execute(Object query) {
        if (!(query instanceof String)) {
            throw new IllegalArgumentException("String query expected");
        }

        PreparedStatement statement = null;
        try {
            statement = connection.prepareStatement((String) query);
            return statement.execute();
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            if (statement != null) {
                try {
                    statement.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
        }

        throw new RuntimeException("Cannot execute query");
    }
}

If connection pool stored in the connection supports connection release, we return SQL connection back to the pool. Otherwise the connection is closed.
execute method checks if query is an instance of String and if yes, uses PreparedStatement to execute query.

Now, let us implement connection pool:

public class ConnectionPoolImpl implements ConnectionPool, ReleaseConnection {
    @Override
    public Connection getConnection() {
        // todo: implement getConnection method
        return null;
    }

    @Override
    public void close(java.sql.Connection connection) {
        // todo: implement close method
    }
}

We will need to pass driverClassName, jdbcUrl, username and password for the pool to be able to create new connections. Since number of connections that the pool can create is not limited, it is easier to use a thread-safe collection – BlockingQueue:

import java.sql.Driver;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;

public class ConnectionPoolImpl implements ConnectionPool, ReleaseConnection {
    private final Driver driver;
    private final String jdbcUrl;
    private final String username;
    private final String password;
    private final BlockingQueue<java.sql.Connection> queue;

    public ConnectionPoolImpl(String driverClassName, String jdbcUrl,
                              String username, String password)
            throws ClassNotFoundException, IllegalAccessException, InstantiationException {
        Class c = Class.forName(driverClassName);
        this.driver = (Driver) c.newInstance();

        this.jdbcUrl = jdbcUrl;
        this.username = username;
        this.password = password;
        this.queue = new LinkedBlockingDeque<>();
    }

    @Override
    public Connection getConnection() {
        // todo: implement getConnection method
        return null;
    }

    @Override
    public void close(java.sql.Connection connection) {
        // todo: implement close method
    }
}

Close method is simple – we just need to return the connection back to the queue:

@Override
public void close(java.sql.Connection connection) {
    queue.offer(connection);
}

Now let us move on to getConnection method:

@Override
public Connection getConnection() {
    java.sql.Connection connection = queue.poll();

    if (connection == null) {
        // todo: create and return new connection
        return null;
    }

    return new ConnectionImpl(this, connection);
}

Here we call method on the queue. poll method is thread safe (LinkedBlockingDeque uses ReentrantLock to ensure thread safety), it either removes first SQL connection from the queue or returns null immediately. If connection is not null, we wrap pass it to ConnectionImpl constructor and return to calling code.

If there are no connections in the pool, we need to create a new one:

private java.sql.Connection createNewConnection() {
    try {
        java.util.Properties info = new java.util.Properties();
        info.put("user", username);
        info.put("password", password);

        return driver.connect(jdbcUrl, info);
    } catch (Throwable t) {
        t.printStackTrace();
        throw new RuntimeException("Connection not available", t);
    }
}

@Override
public Connection getConnection() {
    java.sql.Connection connection = queue.poll();

    if (connection == null) {
        connection = createNewConnection();
    }

    return new ConnectionImpl(this, connection);
}

Final code:
Connection.java:

interface Connection {
    /**
     * Closes the connection or returns it back to the connection pool.
     */
    void close();

    /**
     * Executes the query and returns the result.
     * @param query query to execute.
     * @return result of query execution.
     */
    Object execute (Object query);
}

ConnectionImpl.java:

import java.sql.PreparedStatement;
import java.sql.SQLException;

public class ConnectionImpl implements Connection {
    private final ConnectionPool pool;
    private final java.sql.Connection connection;

    public ConnectionImpl(ConnectionPool pool, java.sql.Connection connection) {
        this.pool = pool;
        this.connection = connection;
    }

    @Override
    public void close() {
        if (pool instanceof ReleaseConnection) {
            ((ReleaseConnection) pool).close(connection);
        } else {
            try {
                connection.close();
            } catch (SQLException e) {
                e.printStackTrace();
            }
        }
    }

    @Override
    public Object execute(Object query) {
        if (!(query instanceof String)) {
            throw new IllegalArgumentException("String query expected");
        }

        PreparedStatement statement = null;
        try {
            statement = connection.prepareStatement((String) query);
            return statement.execute();
        } catch (SQLException e) {
            e.printStackTrace();
        } finally {
            if (statement != null) {
                try {
                    statement.close();
                } catch (SQLException e) {
                    e.printStackTrace();
                }
            }
        }

        throw new RuntimeException("Cannot execute query");
    }
}

ConnectionPool.java:

interface ConnectionPool {
    /**
     * Gets a connection from the pool (if it is available) or creates a new one.
     * @return connection
     */
    Connection getConnection();
}

ConnectionPoolImpl.java:

import java.sql.Driver;
import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingDeque;

public class ConnectionPoolImpl implements ConnectionPool, ReleaseConnection {
    private final Driver driver;
    private final String jdbcUrl;
    private final String username;
    private final String password;
    private final BlockingQueue<java.sql.Connection> queue;

    public ConnectionPoolImpl(String driverClassName, String jdbcUrl,
                              String username, String password)
            throws ClassNotFoundException, IllegalAccessException, InstantiationException {
        Class c = Class.forName(driverClassName);
        this.driver = (Driver) c.newInstance();

        this.jdbcUrl = jdbcUrl;
        this.username = username;
        this.password = password;
        this.queue = new LinkedBlockingDeque<>();
    }

    private java.sql.Connection createNewConnection() {
        try {
            java.util.Properties info = new java.util.Properties();
            info.put("user", username);
            info.put("password", password);

            return driver.connect(jdbcUrl, info);
        } catch (Throwable t) {
            t.printStackTrace();
            throw new RuntimeException("Connection not available", t);
        }
    }

    @Override
    public Connection getConnection() {
        java.sql.Connection connection = queue.poll();

        if (connection == null) {
            connection = createNewConnection();
        }

        return new ConnectionImpl(this, connection);
    }

    @Override
    public void close(java.sql.Connection connection) {
        queue.offer(connection);
    }
}

ReleaseConnection.java:

public interface ReleaseConnection {
    /**
     * Return SQL connection back to the pool.
     * @param connection SQL connection to return
     */
    void close(java.sql.Connection connection);
}

Comments (1)