[utils] Add MutexMap

utils/priority_mutex.go

@@ -31,19 +31,19 @@ type PriorityMutex struct {
 	high []chan struct{}
 	low  []chan struct{}
 
-	m sync.Mutex
-	l bool
+	mutex sync.Mutex
+	lock  bool
 }
 
 // Lock attempts to acquire either a high or low
 // priority mutex. When priority is true, a lock
 // will be granted before other low priority callers.
 func (m *PriorityMutex) Lock(priority bool) {
-	m.m.Lock()
+	m.mutex.Lock()
 
-	if !m.l {
-		m.l = true
-		m.m.Unlock()
+	if !m.lock {
+		m.lock = true
+		m.mutex.Unlock()
 		return
 	}
 
@@ -54,16 +54,16 @@ func (m *PriorityMutex) Lock(priority bool) {
 		m.low = append(m.low, c)
 	}
 
-	m.m.Unlock()
+	m.mutex.Unlock()
 	<-c
 }
 
 // Unlock selects the next highest priority lock
 // to grant. If there are no locks to grant, it
-// sets the value of m.l to false.
+// sets the value of m.lock to false.
 func (m *PriorityMutex) Unlock() {
-	m.m.Lock()
-	defer m.m.Unlock()
+	m.mutex.Lock()
+	defer m.mutex.Unlock()
 
 	if len(m.high) > 0 {
 		c := m.high[0]
@@ -79,8 +79,8 @@ func (m *PriorityMutex) Unlock() {
 		return
 	}
 
-	// We only set m.l to false when there are
+	// We only set m.lock to false when there are
 	// no items to unlock because it could cause
 	// lock contention for the next lock to fetch it.
-	m.l = false
+	m.lock = false
 }

utils/priority_mutex_map.go

@@ -0,0 +1,110 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package utils
+
+import (
+	"sync"
+)
+
+// MutexMap is a struct that allows for
+// acquiring a *PriorityMutex via a string identifier
+// or for acquiring a global mutex that blocks
+// the acquisition of any identifier mutexes.
+//
+// This is useful for coordinating concurrent, non-overlapping
+// writes in the storage package.
+type MutexMap struct {
+	entries     map[string]*mutexMapEntry
+	mutex       sync.Mutex
+	globalMutex sync.RWMutex
+}
+
+// mutexMapEntry is the primitive used
+// to track claimed *PriorityMutex.
+type mutexMapEntry struct {
+	lock  *PriorityMutex
+	count int
+}
+
+// NewMutexMap returns a new *MutexMap.
+func NewMutexMap() *MutexMap {
+	return &MutexMap{
+		entries: map[string]*mutexMapEntry{},
+	}
+}
+
+// GLock acquires an exclusive lock across
+// an entire *MutexMap.
+func (m *MutexMap) GLock() {
+	m.globalMutex.Lock()
+}
+
+// GUnlock releases an exclusive lock
+// held for an entire *MutexMap.
+func (m *MutexMap) GUnlock() {
+	m.globalMutex.Unlock()
+}
+
+// Lock acquires a lock for a particular identifier, as long
+// as no other caller has the global mutex or a lock
+// by the same identifier.
+func (m *MutexMap) Lock(identifier string, priority bool) {
+	// We acquire a RLock on m.globalMutex before
+	// acquiring our identifier lock to ensure no
+	// goroutine holds an identifier mutex while
+	// the m.globalMutex is also held.
+	m.globalMutex.RLock()
+
+	// We acquire m when adding items to m.table
+	// so that we don't accidentally overwrite
+	// lock created by another goroutine.
+	m.mutex.Lock()
+	l, ok := m.entries[identifier]
+	if !ok {
+		l = &mutexMapEntry{
+			lock: new(PriorityMutex),
+		}
+		m.entries[identifier] = l
+	}
+	l.count++
+	m.mutex.Unlock()
+
+	// Once we have a m.globalMutex.RLock, it is
+	// safe to acquire an identifier lock.
+	l.lock.Lock(priority)
+}
+
+// Unlock releases a lock held for a particular identifier.
+func (m *MutexMap) Unlock(identifier string) {
+	// The lock at a particular identifier MUST
+	// exist by the time we unlock, otherwise
+	// it would not have been possible to get
+	// the lock to begin with.
+	m.mutex.Lock()
+	entry := m.entries[identifier]
+	if entry.count <= 1 { // this should never be < 0
+		delete(m.entries, identifier)
+	} else {
+		entry.count--
+		entry.lock.Unlock()
+	}
+	m.mutex.Unlock()
+
+	// We release the globalMutex after unlocking
+	// the identifier lock, otherwise it would be possible
+	// for GLock to be acquired while still holding some
+	// lock in the table.
+	m.globalMutex.RUnlock()
+}

utils/priority_mutex_map_test.go

@@ -0,0 +1,87 @@
+// Copyright 2020 Coinbase, Inc.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package utils
+
+import (
+	"context"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"golang.org/x/sync/errgroup"
+)
+
+func TestMutexMap(t *testing.T) {
+	arr := []string{}
+	m := NewMutexMap()
+	g, _ := errgroup.WithContext(context.Background())
+
+	// Lock while adding all locks
+	m.GLock()
+
+	// Add another GLock
+	g.Go(func() error {
+		m.GLock()
+		arr = append(arr, "global-b")
+		m.GUnlock()
+		return nil
+	})
+
+	// To test locking, we use channels
+	// that will cause deadlock if not executed
+	// concurrently.
+	a := make(chan struct{})
+	b := make(chan struct{})
+
+	g.Go(func() error {
+		m.Lock("a", false)
+		assert.Equal(t, m.entries["a"].count, 1)
+		<-a
+		arr = append(arr, "a")
+		close(b)
+		m.Unlock("a")
+		return nil
+	})
+
+	g.Go(func() error {
+		m.Lock("b", false)
+		assert.Equal(t, m.entries["b"].count, 1)
+		close(a)
+		<-b
+		arr = append(arr, "b")
+		m.Unlock("b")
+		return nil
+	})
+
+	time.Sleep(1 * time.Second)
+
+	// Ensure number of expected locks is correct
+	assert.Len(t, m.entries, 0)
+	arr = append(arr, "global-a")
+	m.GUnlock()
+	assert.NoError(t, g.Wait())
+
+	// Check results array to ensure all of the high priority items processed first,
+	// followed by all of the low priority items.
+	assert.Equal(t, []string{
+		"global-a",
+		"a",
+		"b",
+		"global-b", // must wait until all other locks complete
+	}, arr)
+
+	// Ensure lock is no longer occupied
+	assert.Len(t, m.entries, 0)
+}

utils/priority_mutex_test.go

@@ -69,5 +69,5 @@ func TestPriorityMutex(t *testing.T) {
 	assert.Equal(t, expected, arr)
 
 	// Ensure lock is no longer occupied
-	assert.False(t, l.l)
+	assert.False(t, l.lock)
 }