Laniakea/scene_locks.go

package laniakea

import (
	"sort"
	"sync"
)

type sceneLockEntry struct {
	mu   sync.Mutex
	refs int
}

type sceneKeyLocker struct {
	mu      sync.Mutex
	entries map[string]*sceneLockEntry
}

func (l *sceneKeyLocker) lock(keys []string) func() {
	keys = uniqueSortedStrings(keys)
	if len(keys) == 0 {
		return func() {}
	}

	l.mu.Lock()
	if l.entries == nil {
		l.entries = make(map[string]*sceneLockEntry)
	}
	entries := make([]*sceneLockEntry, len(keys))
	for i, key := range keys {
		entry := l.entries[key]
		if entry == nil {
			entry = new(sceneLockEntry)
			l.entries[key] = entry
		}
		entry.refs++
		entries[i] = entry
	}
	l.mu.Unlock()

	for _, entry := range entries {
		entry.mu.Lock()
	}

	return func() {
		for i := len(entries) - 1; i >= 0; i-- {
			entries[i].mu.Unlock()
		}

		l.mu.Lock()
		for i, key := range keys {
			entries[i].refs--
			if entries[i].refs == 0 {
				delete(l.entries, key)
			}
		}
		l.mu.Unlock()
	}
}

func uniqueSortedStrings(values []string) []string {
	sort.Strings(values)
	result := values[:0]
	for _, value := range values {
		if value == "" || len(result) > 0 && result[len(result)-1] == value {
			continue
		}
		result = append(result, value)
	}
	return result
}