How to Create an Easy Thread-Safe Object Pool in Python with ContextManager

Introduction

Sometimes, to save time and resources, it’s handy to have a collection of ready-to-use items, which we call a “pool.” This is especially helpful for things that are costly to create, like database connections. Here’s the simple concept: You gather these items in a group, often in a list, and have a manager who gives them out when required and keeps an eye on how they’re used.

One problem we face is that, besides getting objects from a pool, we also need to return them. We could do this manually, but it’s more elegant to have it done automatically, and that’s where a Context Manager comes in.

Implementation in Python

We will start by defining our imports:

import threading
import logging
from typing import Self
from collections import deque
from contextlib import contextmanager

Line by line:

1. The threading module provides Python’s support for mulithreading
2. Because we want some logging the logging module is imported
3. The Self in the typing represents the current class type.
4. From the collections module we import the deque, which stands for double-ended queue. This is an efficient data structure for adding and removing elements from both ends of the collection.
5. Finally from the contextlib we import the contextmanager attribute. This enables creating custom context managers using generator function.

Next we need to define the object in our object pool, this case a DbConnection:

class DbConnection:
    def __init__(self, initial_id: int):
        self._id = initial_value

    @property
    def id(self) -> int:
        return self._id

This is very simple class, as its only attribute is an id.

Next we come to the ConnectionPool itself:

This looks a bit complicated, so let’s break it down:

class ConnectionPool:
    def __init__(self,initial_connections:list[DbConnection]):
        self._available_connections = deque(initial_connections)
        self._active_connections = []
        self._lock = threading.Lock()
        self._logger = logging.getLogger(__name__)
    
    def get_connection(self)->DbConnection|None:
        with self._lock:
            if len(self._available_connections)==0:
                self._logger.warning("No connections available in pool")
                return None
            result:DbConnection=self._available_connections.popleft()
            self._active_connections.append(result)
            self._logger.info(f"Connection {result.id} acquired")
            return result

    def remove_from_active(self,connection:DbConnection)->DbConnection|None:
        for conn in self._active_connections:
            if conn.id == connection.id:
                self._active_connections.remove(conn)
                return conn
        return None

    def release_connection(self,connection:DbConnection):
        with self._lock:
            remove_result=self.remove_from_active(connection)
            if remove_result is not None:
                self._available_connections.append(remove_result)

    def _release_all(self):
        with self._lock:
            current_connections= self._active_connections.copy()
            for conn in current_connections:
                remove_result=self.remove_from_active(conn)
                if remove_result is not None:
                    self._available_connections.append(remove_result)


    @contextmanager
    def get_connection_ctx(self):
        """Context manager for automatically releasing connections"""
        connection = self.get_connection()
        if connection is None:
            raise RuntimeError("No connections available")
        try:
            yield connection
        finally:
            self.release_connection(connection)

    @property 
    def available_count(self) -> int:
        """Number of available connections"""
        with self._lock:
            return len(self._available_connections)
    
    @property
    def active_count(self) -> int:
        """Number of active connections"""
        with self._lock:
            return len(self._active_connections)
    
    @property
    def total_count(self) -> int:
        """Total number of connections in pool"""
        with self._lock:
            return len(self._available_connections) + len(self._active_connections)

The constructor

The constructor looks like this:

    def __init__(self,initial_connections:list[DbConnection]):
        self._available_connections = deque(initial_connections)
        self._active_connections = []
        self._lock = threading.Lock()
        self._logger = logging.getLogger(__name__)

In summary, we have a list of free connections _available_connections, a list for connections in use _active_connections, and a lock called _lock to make the code thread-safe. Also, because we want to log some info, we initialize a logger. The constructor receives a list of initial connections and initializes the attributes.

Getting connections

For our connection pool we need to be able to get ready-made connections from the pool:

    def get_connection(self)->DbConnection|None:
        with self._lock:
            if len(self._available_connections)==0:
                self._logger.warning("No connections available in pool")
                return None
            result:DbConnection=self._available_connections.popleft()
            self._active_connections.append(result)
            self._logger.info(f"Connection {result.id} acquired")
            return result

Line by line:

1. The method returns either a DbConnection or None if none is available.
2. We begin by acquiring a thread lock, to prevent race conditions when multiple threads attempt to access the pool simultaneously. Basically: only one thread can access the pool’s state at a time. This helps maintain integrity in concurrent environments.
3. Next check if there are any connections available, if no connections are available return a None value.
4. If there are connections available, use popleft() to efficiently remove the first element from the deque.
5. The connection is immediately added to the list of active connections.
6. A message is logged, and the connection is returned to the caller.

Removing an active connection

There’s also a method for removing an active connection from the pool:

    def remove_from_active(self, connection: DbConnection) -> DbConnection | None:
        for i, conn in enumerate(self._active_connections.copy()):
            if conn.id == connection.id:
                return self._active_connections.pop(i)
        return None

This method removes a connection from the list of active connections. It does this by enumerating over a copy of the active connections.

Releasing a single connection

Returning a single connection to the pool is achieved with this method:

    def release_connection(self,connection:DbConnection):
        with self._lock:
            remove_result=self.remove_from_active(connection)
            if remove_result is not None:
                self._available_connections.append(remove_result)

It locks the pool, removes the connection from the active connections, and adds it back to the idle connections.

Releasing all connections

When exiting the context manager, all connections are returned:

    def _release_all(self):
        with self._lock:
            current_connections= self._active_connections.copy()
            for conn in current_connections:
                remove_result=self.remove_from_active(conn)
                if remove_result is not None:
                    self._available_connections.append(remove_result)

This method acquires a lock on the pool, creates a copy of the active connections list, iterates over them, releases each connection, and adds it back to the idle connections.

The contextmanager method

This method creates a contextmanager using the @contextmanager attribute:

    @contextmanager
    def get_connection_ctx(self):

        connection = self.get_connection()
        if connection is None:
            raise RuntimeError("No connections available")
        try:
            yield connection
        finally:
            self.release_connection(connection)

Line by line:

1. The @contextmanager decorator transforms the this generator function into a proper context, i.e. one which could be used with a with statement.
2. First, we try and get a connection. If no connection can be acquired, a RuntimeError is raised. This may seem a bit much, but callers expect a valid connection, if none can be found the error needs to be handled.
3. If a connection could be acquired, it is yielded to the caller.
4. Once the with statement ends, the finally block is called, which releases the connection.

Remeber the finally block is always executed, so there is always a cleanup.

Some utility methods

Finally some utility methods, these will be used to display the state of the pool, showing the number of active, available and total connections:

    @property 
    def available_count(self) -> int:
        """Number of available connections"""
        with self._lock:
            return len(self._available_connections)
    
    @property
    def active_count(self) -> int:
        """Number of active connections"""
        with self._lock:
            return len(self._active_connections)
    
    @property
    def total_count(self) -> int:
        """Total number of connections in pool"""
        with self._lock:
            return len(self._available_connections) + len(self._active_connections)

Testing time

Now we can test our little setup. We will use two approaches, one with the new get_connection_ctx() method, and the other the classic method:

if __name__ == "__main__":

    logging.basicConfig(level=logging.INFO, format='%(levelname)s: %(message)s')
    
    connections: list[DbConnection] = [DbConnection(1), DbConnection(2)]
    connection_pool: ConnectionPool = ConnectionPool(connections)
    
    print(f"Pool status - Available: {connection_pool.available_count}, Active: {connection_pool.active_count}")
    
    print("\n=== Using context manager for individual connections ===")
    try:
        with connection_pool.get_connection_ctx() as conn1:
            print(f"Got connection with id {conn1.id}")
            print(f"Pool status - Available: {connection_pool.available_count}, Active: {connection_pool.active_count}")
            
            with connection_pool.get_connection_ctx() as conn2:
                print(f"Got connection with id {conn2.id}")
                print(f"Pool status - Available: {connection_pool.available_count}, Active: {connection_pool.active_count}")
        
        print(f"After context managers - Available: {connection_pool.available_count}, Active: {connection_pool.active_count}")
    except RuntimeError as e:
        print(f"Error: {e}")

In the testing section, we create a ConnectionPool with two connections and use the context manager to get and return connections. This demonstrates how the combination of a contextmanager and a connection pool can simplify resource management.

Conclusion

The Object Pool pattern is a useful technique for efficiently managing and reusing resources. Using a contextmanager makes it even more convenient and secure. This is because of the error-handling capabilities, and above all because it enables automatic clean up.