From 3ee85b9ee859bd07b942ea75bc2967598968d651 Mon Sep 17 00:00:00 2001
From: Matthew Rocklin <mrocklin@gmail.com>
Date: Thu, 19 Oct 2017 16:25:59 -0400
Subject: [PATCH] lock.acquire(timeout=) returns True/False if received lock

This matches threading.Lock behavior
---
 distributed/lock.py                  | 16 ++++++++++++++--
 distributed/tests/py3_test_client.py |  4 ++--
 distributed/tests/test_locks.py      | 14 +++++++++-----
 3 files changed, 25 insertions(+), 9 deletions(-)

diff --git a/distributed/lock.py b/distributed/lock.py
index 042bdfa28e8..7dcc9584bd9 100644
--- a/distributed/lock.py
+++ b/distributed/lock.py
@@ -45,11 +45,19 @@ def acquire(self, stream=None, name=None, id=None, timeout=None):
                     future = gen.with_timeout(timedelta(seconds=timeout), future)
                 try:
                     yield future
+                except gen.TimeoutError:
+                    result = False
+                else:
+                    result = True
                 finally:
                     event2 = self.events[name].popleft()
                     assert event is event2
-            assert name not in self.ids
-            self.ids[name] = id
+            else:
+                result = True
+            if result:
+                assert name not in self.ids
+                self.ids[name] = id
+            raise gen.Return(result)
 
     def release(self, stream=None, name=None, id=None):
         with log_errors():
@@ -97,6 +105,10 @@ def acquire(self, timeout=None):
         --------
         >>> lock = Lock('x')  # doctest: +SKIP
         >>> lock.acquire(timeout=1)  # doctest: +SKIP
+
+        Returns
+        -------
+        True or False whether or not it sucessfully acquired the lock
         """
         result = self.client.sync(self.client.scheduler.lock_acquire,
                                   name=self.name, id=self.id, timeout=timeout)
diff --git a/distributed/tests/py3_test_client.py b/distributed/tests/py3_test_client.py
index 6ea2b660a3e..8abd568d385 100644
--- a/distributed/tests/py3_test_client.py
+++ b/distributed/tests/py3_test_client.py
@@ -107,7 +107,7 @@ async def f():
             assert c.asynchronous == True
             async with Lock('x'):
                 lock2 = Lock('x')
-                with pytest.raises(gen.TimeoutError):
-                    await lock2.acquire(timeout=0.1)
+                result = await lock2.acquire(timeout=0.1)
+                assert result is False
 
     loop.run_sync(f)
diff --git a/distributed/tests/test_locks.py b/distributed/tests/test_locks.py
index f8fc791452c..e226afb7e26 100644
--- a/distributed/tests/test_locks.py
+++ b/distributed/tests/test_locks.py
@@ -34,15 +34,20 @@ def f(x):
 def test_timeout(c, s, a, b):
     locks = s.extensions['locks']
     lock = Lock('x')
-    yield lock.acquire()
+    result = yield lock.acquire()
+    assert result is True
+    assert locks.ids['x'] == lock.id
+
     lock2 = Lock('x')
     assert lock.id != lock2.id
 
     start = time()
-    with pytest.raises(gen.TimeoutError):
-        yield lock2.acquire(timeout=0.1)
+    result = yield lock2.acquire(timeout=0.1)
     stop = time()
     assert stop - start < 0.3
+    assert result is False
+    assert locks.ids['x'] == lock.id
+    assert not locks.events['x']
 
     yield lock.release()
 
@@ -64,8 +69,7 @@ def test_timeout_sync(loop):
     with cluster() as (s, [a, b]):
         with Client(s['address'], loop=loop) as c:
             with Lock('x') as lock:
-                with pytest.raises(gen.TimeoutError):
-                    Lock('x').acquire(timeout=0.1)
+                assert Lock('x').acquire(timeout=0.1) is False
 
 
 @gen_cluster(client=True)