Use with-statements to acquire locks.

ajdavis · ajdavis · commit f91c3d61cbc8 · 2014-10-02T10:07:43.000-04:00
We can use with-statements now that we've dropped support
for Python 2.4.
diff --git a/bson/objectid.py b/bson/objectid.py
@@ -158,10 +158,9 @@ def __generate(self):
         oid += struct.pack(">H", os.getpid() % 0xFFFF)
 
         # 3 bytes inc
-        ObjectId._inc_lock.acquire()
-        oid += struct.pack(">i", ObjectId._inc)[1:4]
-        ObjectId._inc = (ObjectId._inc + 1) % 0xFFFFFF
-        ObjectId._inc_lock.release()
+        with ObjectId._inc_lock:
+            oid += struct.pack(">i", ObjectId._inc)[1:4]
+            ObjectId._inc = (ObjectId._inc + 1) % 0xFFFFFF
 
         self.__id = oid
 
diff --git a/pymongo/pool.py b/pymongo/pool.py
@@ -339,15 +339,8 @@ def reset(self):
         self.pool_id += 1
         self.pid = os.getpid()
 
-        sockets = None
-        try:
-            # Swapping variables is not atomic. We need to ensure no other
-            # thread is modifying self.sockets, or replacing it, in this
-            # critical section.
-            self.lock.acquire()
+        with self.lock:
             sockets, self.sockets = self.sockets, set()
-        finally:
-            self.lock.release()
 
         for sock_info in sockets:
             sock_info.close()
@@ -477,15 +470,11 @@ def _get_socket_no_auth(self):
 
         # We've now acquired the semaphore and must release it on error.
         try:
-            sock_info, from_pool = None, None
             try:
-                try:
-                    # set.pop() isn't atomic in Jython less than 2.7, see
-                    # http://bugs.jython.org/issue1854
-                    self.lock.acquire()
+                # set.pop() isn't atomic in Jython less than 2.7, see
+                # http://bugs.jython.org/issue1854
+                with self.lock:
                     sock_info, from_pool = self.sockets.pop(), True
-                finally:
-                    self.lock.release()
             except KeyError:
                 sock_info, from_pool = self.connect(), False
 
@@ -551,8 +540,7 @@ def maybe_return_socket(self, sock_info):
     def _return_socket(self, sock_info):
         """Return socket to the pool. If pool is full the socket is discarded.
         """
-        try:
-            self.lock.acquire()
+        with self.lock:
             max_size = self.opts.max_pool_size
             too_many_sockets = (max_size is not None
                                 and len(self.sockets) >= max_size)
@@ -561,8 +549,6 @@ def _return_socket(self, sock_info):
                 self.sockets.add(sock_info)
             else:
                 sock_info.close()
-        finally:
-            self.lock.release()
 
         self._socket_semaphore.release()
 
diff --git a/pymongo/thread_util.py b/pymongo/thread_util.py
@@ -37,6 +37,12 @@ def acquire(self):
     def release(self):
         pass
 
+    def __enter__(self):
+        pass
+
+    def __exit__(self, t, v, tb):
+        pass
+
 
 class ThreadIdent(object):
     def __init__(self):
@@ -70,13 +76,10 @@ def _make_vigil(self):
         # Threadlocals in Python <= 2.7.0 have race conditions when setting
         # attributes and possibly when getting them, too, leading to weakref
         # callbacks not getting called later.
-        self._lock.acquire()
-        try:
+        with self._lock:
             vigil = getattr(self._local, 'vigil', None)
             if not vigil:
                 self._local.vigil = vigil = ThreadIdent.ThreadVigil()
-        finally:
-            self._lock.release()
 
         return vigil
 
diff --git a/test/mod_wsgi_test/test_client.py b/test/mod_wsgi_test/test_client.py
@@ -104,10 +104,9 @@ def run(self):
 
                 self.errors += 1
 
-            URLGetterThread.counter_lock.acquire()
-            URLGetterThread.counter += 1
-            counter = URLGetterThread.counter
-            URLGetterThread.counter_lock.release()
+            with URLGetterThread.counter_lock:
+                URLGetterThread.counter += 1
+                counter = URLGetterThread.counter
 
             should_print = options.verbose and not counter % 1000
 
