go.temporal.io/server@v1.23.0/common/locks/id_mutex.go (about)

     1  // The MIT License
     2  //
     3  // Copyright (c) 2020 Temporal Technologies Inc.  All rights reserved.
     4  //
     5  // Copyright (c) 2020 Uber Technologies, Inc.
     6  //
     7  // Permission is hereby granted, free of charge, to any person obtaining a copy
     8  // of this software and associated documentation files (the "Software"), to deal
     9  // in the Software without restriction, including without limitation the rights
    10  // to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
    11  // copies of the Software, and to permit persons to whom the Software is
    12  // furnished to do so, subject to the following conditions:
    13  //
    14  // The above copyright notice and this permission notice shall be included in
    15  // all copies or substantial portions of the Software.
    16  //
    17  // THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
    18  // IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
    19  // FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
    20  // AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
    21  // LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
    22  // OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
    23  // THE SOFTWARE.
    24  
    25  package locks
    26  
    27  import (
    28  	"sync"
    29  )
    30  
    31  type (
    32  	// HashFunc represents a hash function for string
    33  	HashFunc func(interface{}) uint32
    34  
    35  	// IDMutex is an interface which can lock on specific comparable identifier
    36  	IDMutex interface {
    37  		LockID(identifier interface{})
    38  		UnlockID(identifier interface{})
    39  	}
    40  
    41  	// idMutexShardImpl is the implementation of IDMutex shard
    42  	idMutexImpl struct {
    43  		numShard uint32
    44  		hashFn   HashFunc
    45  		shards   map[uint32]*idMutexShardImpl
    46  	}
    47  
    48  	// idMutexShardImpl is the implementation of IDMutex shard
    49  	idMutexShardImpl struct {
    50  		sync.Mutex
    51  		mutexInfos map[interface{}]*mutexInfo
    52  	}
    53  
    54  	mutexInfo struct {
    55  		// how many caller are using this lock info, including the
    56  		// the caller already have the lock
    57  		// this is guarded by lock in idLockImpl
    58  		waitCount int
    59  
    60  		// actual lock
    61  		sync.Mutex
    62  	}
    63  )
    64  
    65  // NewIDMutex create a new IDLock
    66  func NewIDMutex(numShard uint32, hashFn HashFunc) IDMutex {
    67  	impl := &idMutexImpl{
    68  		numShard: numShard,
    69  		hashFn:   hashFn,
    70  		shards:   make(map[uint32]*idMutexShardImpl),
    71  	}
    72  	for i := uint32(0); i < numShard; i++ {
    73  		impl.shards[i] = &idMutexShardImpl{
    74  			mutexInfos: make(map[interface{}]*mutexInfo),
    75  		}
    76  	}
    77  
    78  	return impl
    79  }
    80  
    81  func newMutexInfo() *mutexInfo {
    82  	return &mutexInfo{
    83  		waitCount: 1,
    84  	}
    85  }
    86  
    87  // LockID lock by specific identifier
    88  func (idMutex *idMutexImpl) LockID(identifier interface{}) {
    89  	shard := idMutex.shards[idMutex.getShardIndex(identifier)]
    90  
    91  	shard.Lock()
    92  	mutexInfo, ok := shard.mutexInfos[identifier]
    93  	if !ok {
    94  		mutexInfo := newMutexInfo()
    95  		shard.mutexInfos[identifier] = mutexInfo
    96  		shard.Unlock()
    97  		mutexInfo.Lock()
    98  		return
    99  	}
   100  
   101  	mutexInfo.waitCount++
   102  	shard.Unlock()
   103  	mutexInfo.Lock()
   104  }
   105  
   106  // UnlockID unlock by specific identifier
   107  func (idMutex *idMutexImpl) UnlockID(identifier interface{}) {
   108  	shard := idMutex.shards[idMutex.getShardIndex(identifier)]
   109  
   110  	shard.Lock()
   111  	defer shard.Unlock()
   112  	mutexInfo, ok := shard.mutexInfos[identifier]
   113  	if !ok {
   114  		panic("cannot find workflow lock")
   115  	}
   116  	mutexInfo.Unlock()
   117  	if mutexInfo.waitCount == 1 {
   118  		delete(shard.mutexInfos, identifier)
   119  	} else {
   120  		mutexInfo.waitCount--
   121  	}
   122  }
   123  
   124  func (idMutex *idMutexImpl) getShardIndex(key interface{}) uint32 {
   125  	return idMutex.hashFn(key) % idMutex.numShard
   126  }