Add support for lock service API (#50)

* Add support for lock service API

https://coreos.com/etcd/docs/latest/dev-guide/api_concurrency_reference_v3.html

* add timeout to lock service API tests

Author: dariko

etcd3/apis/__init__.py

@@ -5,6 +5,7 @@
 from .extra import ExtraAPI
 from .kv import KVAPI
 from .lease import LeaseAPI
+from .lock import LockAPI
 from .maintenance import MaintenanceAPI
 from .watch import WatchAPI
 
@@ -17,4 +18,5 @@
     'MaintenanceAPI',
     'LeaseAPI',
     'BaseAPI'
+    'LockAPI'
 ]

etcd3/apis/lock.py

@@ -0,0 +1,46 @@
+from .base import BaseAPI
+
+
+class LockAPI(BaseAPI):
+    def lock(self, name, lease=0):
+        """
+        Lock acquires a distributed shared lock on a given named lock.
+        On success, it will return a unique key that exists so long as
+        the lock is held by the caller. This key can be used in
+        conjunction with transactions to safely ensure updates to etcd
+        only occur while holding lock ownership. The lock is held until
+        Unlock is called on the key or the lease associate with the
+        owner expires.
+
+        :type name: str
+        :param name: name is the identifier for the distributed shared lock to be acquired.
+        :type lease: int
+        :param lease: lease is the lease ID to associate with the key in the key-value store. A lease
+            value of 0 indicates no lease.
+        """
+
+        method = '/v3beta/lock/lock'
+        data = {
+            "name": name,
+            "lease": lease
+        }
+        return self.call_rpc(method, data=data)
+
+    def unlock(self, key):
+        """
+        Unlock takes a key returned by Lock and releases the hold on
+        lock. The next Lock caller waiting for the lock will then be
+        woken up and given ownership of the lock.
+
+        :type key: str
+        :param key: key is the lock ownership key granted by Lock.
+        :type lease: int
+        :param lease: lease is the lease ID to associate with the key in the key-value store. A lease
+            value of 0 indicates no lease.
+        """
+
+        method = '/v3beta/lock/unlock'
+        data = {
+            "key": key,
+        }
+        return self.call_rpc(method, data=data)

etcd3/baseclient.py

@@ -14,6 +14,7 @@
 from .apis import ExtraAPI
 from .apis import KVAPI
 from .apis import LeaseAPI
+from .apis import LockAPI
 from .apis import MaintenanceAPI
 from .apis import WatchAPI
 from .stateful import Lease
@@ -46,7 +47,8 @@ def __iter__(self):
         raise NotImplementedError
 
 
-class BaseClient(AuthAPI, ClusterAPI, KVAPI, LeaseAPI, MaintenanceAPI, WatchAPI, ExtraAPI):
+class BaseClient(AuthAPI, ClusterAPI, KVAPI, LeaseAPI, MaintenanceAPI,
+                 WatchAPI, ExtraAPI, LockAPI):
     def __init__(self, host='localhost', port=2379, protocol='http',
                  cert=(), verify=None,
                  timeout=None, headers=None, user_agent=None, pool_size=30,

etcd3/rpc.swagger.json

@@ -799,6 +799,60 @@
         ]
       }
     },
+    "/v3beta/lock/lock": {
+      "post": {
+        "summary": "Lock acquires a distributed shared lock on a given named lock.\nOn success, it will return a unique key that exists so long asnthe lock is held by the caller. This key can be used in\nconjunction with transactions to safely ensure updates to etcdnonly occur while holding lock ownership. The lock is held until\nUnlock is called on the key or the lease associate with thenowner expires.",
+        "operationId": "Lock",
+        "responses": {
+          "200": {
+            "description": "",
+            "schema": {
+              "$ref": "#/definitions/etcdserverpbLockResponse"
+            }
+          }
+        },
+        "parameters": [
+          {
+            "name": "body",
+            "in": "body",
+            "required": true,
+            "schema": {
+              "$ref": "#/definitions/etcdserverpbLockRequest"
+            }
+          }
+        ],
+        "tags": [
+          "Lock"
+        ]
+      }
+    },
+    "/v3beta/lock/unlock": {
+      "post": {
+        "summary": "Unlock takes a key returned by Lock and releases the hold on\nlock. The next Lock caller waiting for the lock will then benwoken up and given ownership of the lock.",
+        "operationId": "Unlock",
+        "responses": {
+          "200": {
+            "description": "",
+            "schema": {
+              "$ref": "#/definitions/etcdserverpbUnlockResponse"
+            }
+          }
+        },
+        "parameters": [
+          {
+            "name": "body",
+            "in": "body",
+            "required": true,
+            "schema": {
+              "$ref": "#/definitions/etcdserverpbUnlockRequest"
+            }
+          }
+        ],
+        "tags": [
+          "Lock"
+        ]
+      }
+    },
     "/v3alpha/maintenance/alarm": {
       "post": {
         "summary": "Alarm activates, deactivates, and queries alarms regarding cluster health.",
@@ -1664,6 +1718,52 @@
         }
       }
     },
+    "etcdserverpbLockRequest": {
+      "type": "object",
+      "properties": {
+        "name": {
+          "type": "string",
+          "format": "byte",
+          "description": "name is the identifier for the distributed shared lock to be acquired."
+        },
+        "lease": {
+          "type": "string",
+          "format": "int64",
+          "description": "lease is the ID of the lease that will be attached to ownership of the lock. If the lease expires or is revoked and currently holds the lock, the lock is automatically released. Calls to Lock with the same lease will be treated as a single acquisition; locking twice with the same lease is a no-op."
+        }
+      }
+    },
+    "etcdserverpbLockResponse": {
+      "type": "object",
+      "properties": {
+        "header": {
+          "$ref": "#/definitions/etcdserverpbResponseHeader"
+        },
+        "key": {
+          "type": "string",
+          "format": "byte",
+          "description": "key is a key that will exist on etcd for the duration that the Lock caller owns the lock. Users should not modify this key or the lock may exhibit undefined behavior."
+        }
+      }
+    },
+    "etcdserverpbUnlockRequest": {
+      "type": "object",
+      "properties": {
+        "key": {
+          "type": "string",
+          "format": "byte",
+          "description": "key is the lock ownership key granted by Lock."
+        }
+      }
+    },
+    "etcdserverpbUnlockResponse": {
+      "type": "object",
+      "properties": {
+        "header": {
+          "$ref": "#/definitions/etcdserverpbResponseHeader"
+        }
+      }
+    },
     "etcdserverpbMember": {
       "type": "object",
       "properties": {

tests/test_lock_apis.py

@@ -0,0 +1,48 @@
+import time
+
+import pytest
+
+from etcd3.client import Client
+from tests.docker_cli import docker_run_etcd_main
+from .envs import protocol, host, port
+from .etcd_go_cli import NO_ETCD_SERVICE, etcdctl
+
+
+@pytest.fixture(scope='module')
+def client():
+    """
+    init Etcd3Client, close its connection-pool when teardown
+    """
+    _, p, _ = docker_run_etcd_main()
+    c = Client(host, p, protocol)
+    yield c
+    c.close()
+
+
+@pytest.mark.skipif(NO_ETCD_SERVICE, reason="no etcd service available")
+def test_hash(client):
+    assert client.hash().hash
+
+@pytest.mark.timeout(60)
+@pytest.mark.skipif(NO_ETCD_SERVICE, reason="no etcd service available")
+def test_lock_flow(client):
+    lease1 = client.Lease(5)
+    lease1.grant()
+    lock1 = client.lock('test_lock', lease1._ID)
+    assert lock1.key.startswith(b'test_lock/')
+
+    lease2 = client.Lease(15)
+    lease2.grant()
+    start_lock_ts = time.time()
+    client.lock('test_lock', lease2._ID)
+    assert (time.time() - start_lock_ts) > 3
+
+    lease2.revoke()
+
+    lease3 = client.Lease(5)
+    lease3.grant()
+    start_lock_ts = time.time()
+    lock3 = client.lock('test_lock', lease3._ID)
+    assert (time.time() - start_lock_ts) < 2
+
+    client.unlock(lock3.key)