github.com/pachyderm/pachyderm@v1.13.4/src/server/pkg/dlock/dlock.go (about)

     1  // Package dlock implements a distributed lock on top of etcd.
     2  package dlock
     3  
     4  import (
     5  	"context"
     6  
     7  	etcd "github.com/coreos/etcd/clientv3"
     8  	"github.com/coreos/etcd/clientv3/concurrency"
     9  )
    10  
    11  // DLock is a handle to a distributed lock.
    12  type DLock interface {
    13  	// Lock acquries the distributed lock, blocking if necessary.  If
    14  	// the lock is acquired, it returns a context that should be used
    15  	// in any subsequent blocking requests, so that if you lose the lock,
    16  	// the requests get cancelled correctly.
    17  	Lock(context.Context) (context.Context, error)
    18  	// Unlock releases the distributed lock.
    19  	Unlock(context.Context) error
    20  }
    21  
    22  type etcdImpl struct {
    23  	client *etcd.Client
    24  	prefix string
    25  
    26  	session *concurrency.Session
    27  	mutex   *concurrency.Mutex
    28  }
    29  
    30  // NewDLock attempts to acquire a distributed lock that locks a given prefix
    31  // in the data store.
    32  func NewDLock(client *etcd.Client, prefix string) DLock {
    33  	return &etcdImpl{
    34  		client: client,
    35  		prefix: prefix,
    36  	}
    37  }
    38  
    39  func (d *etcdImpl) Lock(ctx context.Context) (context.Context, error) {
    40  	// The default TTL is 60 secs which means that if a node dies, it
    41  	// still holds the lock for 60 secs, which is too high.
    42  	session, err := concurrency.NewSession(d.client, concurrency.WithContext(ctx), concurrency.WithTTL(15))
    43  	if err != nil {
    44  		return nil, err
    45  	}
    46  
    47  	mutex := concurrency.NewMutex(session, d.prefix)
    48  	if err := mutex.Lock(ctx); err != nil {
    49  		return nil, err
    50  	}
    51  
    52  	ctx, cancel := context.WithCancel(ctx)
    53  	go func() {
    54  		select {
    55  		case <-ctx.Done():
    56  		case <-session.Done():
    57  			cancel()
    58  		}
    59  	}()
    60  
    61  	d.session = session
    62  	d.mutex = mutex
    63  	return ctx, nil
    64  }
    65  
    66  func (d *etcdImpl) Unlock(ctx context.Context) error {
    67  	if err := d.mutex.Unlock(ctx); err != nil {
    68  		return err
    69  	}
    70  	return d.session.Close()
    71  }