github.com/hashicorp/vault/sdk@v0.13.0/helper/locksutil/locks.go (about)

     1  // Copyright (c) HashiCorp, Inc.
     2  // SPDX-License-Identifier: MPL-2.0
     3  
     4  package locksutil
     5  
     6  import (
     7  	"sync"
     8  
     9  	"github.com/hashicorp/vault/sdk/helper/cryptoutil"
    10  	"github.com/sasha-s/go-deadlock"
    11  )
    12  
    13  const (
    14  	LockCount = 256
    15  )
    16  
    17  // DeadlockRWMutex is the RW version of DeadlockMutex.
    18  type DeadlockRWMutex struct {
    19  	deadlock.RWMutex
    20  }
    21  
    22  type LockEntry struct {
    23  	sync.RWMutex
    24  }
    25  
    26  // CreateLocks returns an array so that the locks can be iterated over in
    27  // order.
    28  //
    29  // This is only threadsafe if a process is using a single lock, or iterating
    30  // over the entire lock slice in order. Using a consistent order avoids
    31  // deadlocks because you can never have the following:
    32  //
    33  // Lock A, Lock B
    34  // Lock B, Lock A
    35  //
    36  // Where process 1 is now deadlocked trying to lock B, and process 2 deadlocked trying to lock A
    37  func CreateLocks() []*LockEntry {
    38  	ret := make([]*LockEntry, LockCount)
    39  	for i := range ret {
    40  		ret[i] = new(LockEntry)
    41  	}
    42  	return ret
    43  }
    44  
    45  func CreateLocksWithDeadlockDetection() []*DeadlockRWMutex {
    46  	ret := make([]*DeadlockRWMutex, LockCount)
    47  	for i := range ret {
    48  		ret[i] = new(DeadlockRWMutex)
    49  	}
    50  	return ret
    51  }
    52  
    53  func LockIndexForKey(key string) uint8 {
    54  	return uint8(cryptoutil.Blake2b256Hash(key)[0])
    55  }
    56  
    57  func LockForKey(locks []*LockEntry, key string) *LockEntry {
    58  	return locks[LockIndexForKey(key)]
    59  }
    60  
    61  func LocksForKeys(locks []*LockEntry, keys []string) []*LockEntry {
    62  	lockIndexes := make(map[uint8]struct{}, len(keys))
    63  	for _, k := range keys {
    64  		lockIndexes[LockIndexForKey(k)] = struct{}{}
    65  	}
    66  
    67  	locksToReturn := make([]*LockEntry, 0, len(keys))
    68  	for i, l := range locks {
    69  		if _, ok := lockIndexes[uint8(i)]; ok {
    70  			locksToReturn = append(locksToReturn, l)
    71  		}
    72  	}
    73  
    74  	return locksToReturn
    75  }
    76  
    77  func LockForKeyWithDeadLockDetection(locks []*DeadlockRWMutex, key string) *DeadlockRWMutex {
    78  	return locks[LockIndexForKey(key)]
    79  }
    80  
    81  func LocksForKeysWithDeadLockDetection(locks []*DeadlockRWMutex, keys []string) []*DeadlockRWMutex {
    82  	lockIndexes := make(map[uint8]struct{}, len(keys))
    83  	for _, k := range keys {
    84  		lockIndexes[LockIndexForKey(k)] = struct{}{}
    85  	}
    86  
    87  	locksToReturn := make([]*DeadlockRWMutex, 0, len(keys))
    88  	for i, l := range locks {
    89  		if _, ok := lockIndexes[uint8(i)]; ok {
    90  			locksToReturn = append(locksToReturn, l)
    91  		}
    92  	}
    93  
    94  	return locksToReturn
    95  }