Ionrock Dot Org

My Weblog

A Contextmanager Based Connection Pool

At work we have a whole variety of services we use that utilize some concept of a connection. Some of the services are RESTful, while others utilize Pyro and straight sockets. Seeing as these are so common, our makeshift web framework has a simple pool implementation that allows you to reuse the connections in a threadsafe way. One of these services is called Faststore (and yes this name is pretty bad). Faststore is a storage tool that aims at making writes extremely fast. It uses a bsddb underneath and handles a massive amount of data for us already. We've also written a CouchDB like app on top of Faststore called Ottoman (which I thought was a very clever name). In both cases, there is a pool that can be made available via CherryPy tools that allows you to use and reuse connections to these data storage services.

Seeing as our goal is to eventually make these apps open source, I've started playing around with them in my free time to see what it would take to remove a few of the more coupled aspects, which lead me to needing a pool implementation. I could have simply taken the pool from our framework, but it seemed like a good opportunity to learn something new and I read this article  from Jesse Noller on context managers that seemed applicable. The result is a simple connection pool using context managers. I call it "poodle" because it kind of sounds like "pool".

from __future__ import with_statement import thread, threading import contextlib import random import time class Pool(object):     def __init__(self, factory, args=None, kwargs=None, cleanup=None, min=5):         self.pool = []         self.swimmers = {}         self.args = args or tuple()         self.kwargs = kwargs or {}         self.factory = factory         self.cleanup = cleanup         self._lock = threading.Semaphore()         for i in xrange(min):             self._ci()             self.pool.append(self._create())     def _create(self):         return self.factory(*self.args, **self.kwargs)     def _get(self):             return self.swimmers[id]     @contextlib.contextmanager     def get(self):         with self._lock:             id = thread.get_ident()             if id not in self.swimmers:                 if self.pool:                     self.swimmers[id] = self.pool.pop()                 else:                     self.swimmers[id] = self._create()             yield self.swimmers[id]         with self._lock:             id = thread.get_ident()             if self.swimmers.get(id):                 if self.cleanup:                     self.cleanup(self.swimmers[id])                 else:                     self.pool.append(self.swimmers[id])                     del self.swimmers[id] class MockThread(threading.Thread):     def __init__(self, pool, name, indent=0):         threading.Thread.__init__(self)         self.pool = pool         self.name = name         self.indent = '\t'.join(['|' for i in xrange(0, indent)])     def m(self, *s):         print '%s %s' % (str(self.indent), ''.join(map(str, s)))     def run(self):         with self.pool.get() as conn:             waiting = random.randint(1, 3)             self.m('got connection')             self.m('using connection in ', self.name)             time.sleep(waiting)             conn(self.indent, ' hello world')         return class MockConn(object):     def __init__(self, name):         self.name = name         def __call__(self, *args):         print ''.join(args) if __name__ == '__main__':     tp = Pool(MockConn, ['eric'], min=0)     workers = [MockThread(tp, x, x) for x in range(0, 10)]     for i, w in enumerate(workers):         w.start()

The win here is that by using context managers, I've eliminated the need to check to see if threads are using a connection. This is traditionally necessary because an exception somewhere can be missed, which in turn never releases the lock on the connection. The context managers should automatically release the lock no matter what, which simplifies the code quite a bit. That said, I have no idea if there are blocking issues I'm missing or anything so feel free to comment and set me straight!

Share: