【代码】基于etcd的分布式锁（python版）

参考k8s的leader选举： k8s.io/client-go/tools/leaderelection

import os
import socket
import time
from threading import Thread

import etcd3


class DistributedLock(object):
    def __init__(self, etcd, name, ttl, renew):
        self.etcd: etcd3.Etcd3Client = etcd
        self.key = f"/locks/{name}"
        self.value = f"{socket.gethostname()}-{os.getpid()}"
        self.renew_value = lambda: self.value + f" {time.time()}"
        self.ttl = ttl
        self.renew = renew
        self.acquired = False
        self.holder = None

    def acquire(self):
        value = self.etcd.get(self.key)[0]
        if value is None:
            print('init acquire')
            succeeded, resp = self.etcd.transaction(compare=[etcd.transactions.create(self.key) == 0],
                                                    success=[etcd.transactions.put(self.key, self.renew_value())],
                                                    failure=[etcd.transactions.get(self.key)])
            if not succeeded:
                print('init acquire failure')
                holder = resp[0][0][0].decode().split(' ')[0]
            else:
                print('init acquire success')
                holder = self.value
        else:
            holder, alive_time = value.decode().split(' ')
            print('current holder: ', holder)
            if float(alive_time) + self.ttl < time.time():
                print('holder may dead, try to acquire')
                succeeded, resp = self.etcd.transaction(compare=[etcd.transactions.value(self.key) == value],
                                                        success=[etcd.transactions.put(self.key, self.renew_value())],
                                                        failure=[etcd.transactions.get(self.key)])
                if not succeeded:
                    print('try acquire failure')
                    holder = resp[0][0][0].decode().split(' ')[0]
                else:
                    print('try acquire success')
                    holder = self.value
            else:
                if holder == self.value:
                    succeeded, resp = self.etcd.transaction(compare=[etcd.transactions.value(self.key) == value],
                                                            success=[
                                                                etcd.transactions.put(self.key, self.renew_value()), ],
                                                            failure=[etcd.transactions.get(self.key)])
                    if succeeded:
                        print('holder is alive, refreshed success')
                    else:
                        print('holder is alive, refreshed failure')
                else:
                    print("holder is alive, can't acquire")
                succeeded = holder == self.value
        print(self.value, "acquired: ", succeeded)
        self.acquired = succeeded
        self.holder = holder

    def release(self):
        value = self.etcd.get(self.key)[0]
        if value is None:
            return
        holder = value.decode().split(' ')
        if holder != self.value:
            return
        succeeded, _ = self.etcd.transaction(compare=[etcd.transactions.value(self.key) == value[0]],
                                             success=[etcd.transactions.put(self.key, self.value)],
                                             failure=[etcd.transactions.get(self.key)])

    def refresh(self):
        def try_acquire_or_refresh():
            while 1:
                self.acquire()
                time.sleep(self.renew)

        thread = Thread(target=try_acquire_or_refresh)
        thread.start()
        thread.join()

    def is_acquired(self):
        return self.acquired


if __name__ == '__main__':
    etcd = etcd3.client()
    lock = DistributedLock(etcd, 'global', 10, 3)
    lock.acquire()
    lock.refresh()