locker.go 2.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112
  1. /*
  2. Package locker provides a mechanism for creating finer-grained locking to help
  3. free up more global locks to handle other tasks.
  4. The implementation looks close to a sync.Mutex, however the user must provide a
  5. reference to use to refer to the underlying lock when locking and unlocking,
  6. and unlock may generate an error.
  7. If a lock with a given name does not exist when `Lock` is called, one is
  8. created.
  9. Lock references are automatically cleaned up on `Unlock` if nothing else is
  10. waiting for the lock.
  11. */
  12. package locker
  13. import (
  14. "errors"
  15. "sync"
  16. "sync/atomic"
  17. )
  18. // ErrNoSuchLock is returned when the requested lock does not exist
  19. var ErrNoSuchLock = errors.New("no such lock")
  20. // Locker provides a locking mechanism based on the passed in reference name
  21. type Locker struct {
  22. mu sync.Mutex
  23. locks map[string]*lockCtr
  24. }
  25. // lockCtr is used by Locker to represent a lock with a given name.
  26. type lockCtr struct {
  27. mu sync.Mutex
  28. // waiters is the number of waiters waiting to acquire the lock
  29. // this is int32 instead of uint32 so we can add `-1` in `dec()`
  30. waiters int32
  31. }
  32. // inc increments the number of waiters waiting for the lock
  33. func (l *lockCtr) inc() {
  34. atomic.AddInt32(&l.waiters, 1)
  35. }
  36. // dec decrements the number of waiters waiting on the lock
  37. func (l *lockCtr) dec() {
  38. atomic.AddInt32(&l.waiters, -1)
  39. }
  40. // count gets the current number of waiters
  41. func (l *lockCtr) count() int32 {
  42. return atomic.LoadInt32(&l.waiters)
  43. }
  44. // Lock locks the mutex
  45. func (l *lockCtr) Lock() {
  46. l.mu.Lock()
  47. }
  48. // Unlock unlocks the mutex
  49. func (l *lockCtr) Unlock() {
  50. l.mu.Unlock()
  51. }
  52. // New creates a new Locker
  53. func New() *Locker {
  54. return &Locker{
  55. locks: make(map[string]*lockCtr),
  56. }
  57. }
  58. // Lock locks a mutex with the given name. If it doesn't exist, one is created
  59. func (l *Locker) Lock(name string) {
  60. l.mu.Lock()
  61. if l.locks == nil {
  62. l.locks = make(map[string]*lockCtr)
  63. }
  64. nameLock, exists := l.locks[name]
  65. if !exists {
  66. nameLock = &lockCtr{}
  67. l.locks[name] = nameLock
  68. }
  69. // increment the nameLock waiters while inside the main mutex
  70. // this makes sure that the lock isn't deleted if `Lock` and `Unlock` are called concurrently
  71. nameLock.inc()
  72. l.mu.Unlock()
  73. // Lock the nameLock outside the main mutex so we don't block other operations
  74. // once locked then we can decrement the number of waiters for this lock
  75. nameLock.Lock()
  76. nameLock.dec()
  77. }
  78. // Unlock unlocks the mutex with the given name
  79. // If the given lock is not being waited on by any other callers, it is deleted
  80. func (l *Locker) Unlock(name string) error {
  81. l.mu.Lock()
  82. nameLock, exists := l.locks[name]
  83. if !exists {
  84. l.mu.Unlock()
  85. return ErrNoSuchLock
  86. }
  87. if nameLock.count() == 0 {
  88. delete(l.locks, name)
  89. }
  90. nameLock.Unlock()
  91. l.mu.Unlock()
  92. return nil
  93. }