github.com/quay/claircore@v1.5.28/libvuln/updates/locks.go (about)

     1  package updates
     2  
     3  import (
     4  	"context"
     5  	"sync"
     6  )
     7  
     8  var _ LockSource = (*localLockSource)(nil)
     9  
    10  type localLockSource struct {
    11  	sync.Mutex
    12  	wait *sync.Cond
    13  	m    map[string]struct{}
    14  }
    15  
    16  // NewLocalLockSource provides locks backed by local concurrency primitives.
    17  func NewLocalLockSource() *localLockSource {
    18  	l := &localLockSource{
    19  		m: make(map[string]struct{}),
    20  	}
    21  	l.wait = sync.NewCond(&l.Mutex)
    22  	return l
    23  }
    24  
    25  // BUG(hank) The API provided by localLockSource does not respect the parent
    26  // context cancellation when waiting for a lock.
    27  
    28  func (s *localLockSource) Lock(ctx context.Context, key string) (context.Context, context.CancelFunc) {
    29  	s.Mutex.Lock()
    30  	defer s.Mutex.Unlock()
    31  	for _, exists := s.m[key]; exists; _, exists = s.m[key] {
    32  		s.wait.Wait()
    33  	}
    34  	s.m[key] = struct{}{}
    35  	c, f := context.WithCancel(ctx)
    36  	return c, s.cancelfunc(key, f)
    37  }
    38  
    39  func (s *localLockSource) TryLock(ctx context.Context, key string) (context.Context, context.CancelFunc) {
    40  	c, f := context.WithCancel(ctx)
    41  	s.Mutex.Lock()
    42  	defer s.Mutex.Unlock()
    43  	if _, exists := s.m[key]; exists {
    44  		f()
    45  		return c, f
    46  	}
    47  	s.m[key] = struct{}{}
    48  	return c, s.cancelfunc(key, f)
    49  }
    50  
    51  func (s *localLockSource) Close(ctx context.Context) error {
    52  	return nil
    53  }
    54  
    55  // Cancelfunc returns a CancelFunc that calls "next" and then unlocks.
    56  func (s *localLockSource) cancelfunc(key string, next context.CancelFunc) context.CancelFunc {
    57  	return func() {
    58  		next()
    59  		s.Mutex.Lock()
    60  		defer s.Mutex.Unlock()
    61  		delete(s.m, key)
    62  		s.wait.Broadcast()
    63  	}
    64  }
    65  
    66  // Peek reports whether the key is locked or not.
    67  func (s *localLockSource) peek(key string) bool {
    68  	s.Mutex.Lock()
    69  	defer s.Mutex.Unlock()
    70  	_, ok := s.m[key]
    71  	return ok
    72  }