feat: add RWKeyedLocker

This commit is contained in:
dudaodong
2025-04-17 14:25:19 +08:00
parent f4dee28ebb
commit d7f3354b98
4 changed files with 232 additions and 61 deletions

View File

@@ -58,6 +58,7 @@ func (l *KeyedLocker[K]) Do(ctx context.Context, key K, fn func()) error {
}
}
// acquire tries to acquire a lock for the specified key.
func (l *KeyedLocker[K]) acquire(key K) *lockEntry {
lock, _ := l.locks.LoadOrStore(key, &lockEntry{})
entry := lock.(*lockEntry)
@@ -70,6 +71,7 @@ func (l *KeyedLocker[K]) acquire(key K) *lockEntry {
return entry
}
// release releases the lock for the specified key.
func (l *KeyedLocker[K]) release(key K, entry *lockEntry, rawKey K) {
if atomic.AddInt32(&entry.ref, -1) == 0 {
entry.mu.Lock()
@@ -91,3 +93,101 @@ func (l *KeyedLocker[K]) release(key K, entry *lockEntry, rawKey K) {
}
}
}
// RWKeyedLocker is a read-write version of KeyedLocker.
type RWKeyedLocker[K comparable] struct {
locks sync.Map
ttl time.Duration
}
type rwLockEntry struct {
mu sync.RWMutex
ref int32
timer atomic.Pointer[time.Timer]
}
// NewRWKeyedLocker creates a new RWKeyedLocker with the specified TTL for lock expiration.
// The TTL is used to automatically release locks that are no longer held.
func NewRWKeyedLocker[K comparable](ttl time.Duration) *RWKeyedLocker[K] {
return &RWKeyedLocker[K]{ttl: ttl}
}
// RLock acquires a read lock for the specified key and executes the provided function.
// It returns an error if the context is canceled before the function completes.
func (l *RWKeyedLocker[K]) RLock(ctx context.Context, key K, fn func()) error {
entry := l.acquire(key)
defer l.release(entry, key)
done := make(chan struct{})
go func() {
entry.mu.RLock()
defer entry.mu.RUnlock()
select {
case <-ctx.Done():
default:
fn()
}
close(done)
}()
select {
case <-ctx.Done():
return ctx.Err()
case <-done:
return nil
}
}
// Lock acquires a write lock for the specified key and executes the provided function.
// It returns an error if the context is canceled before the function completes.
func (l *RWKeyedLocker[K]) Lock(ctx context.Context, key K, fn func()) error {
entry := l.acquire(key)
defer l.release(entry, key)
done := make(chan struct{})
go func() {
entry.mu.Lock()
defer entry.mu.Unlock()
select {
case <-ctx.Done():
default:
fn()
}
close(done)
}()
select {
case <-ctx.Done():
return ctx.Err()
case <-done:
return nil
}
}
// acquire tries to acquire a read lock for the specified key.
func (l *RWKeyedLocker[K]) acquire(key K) *rwLockEntry {
actual, _ := l.locks.LoadOrStore(key, &rwLockEntry{})
entry := actual.(*rwLockEntry)
atomic.AddInt32(&entry.ref, 1)
if t := entry.timer.Swap(nil); t != nil {
t.Stop()
}
return entry
}
// release releases the lock for the specified key.
func (l *RWKeyedLocker[K]) release(entry *rwLockEntry, rawKey K) {
if atomic.AddInt32(&entry.ref, -1) == 0 {
timer := time.AfterFunc(l.ttl, func() {
if atomic.LoadInt32(&entry.ref) == 0 {
l.locks.Delete(rawKey)
}
})
entry.timer.Store(timer)
}
}

View File

@@ -104,3 +104,74 @@ func TestKeyedLocker_LockReleaseAfterTTL(t *testing.T) {
err = locker.Do(context.Background(), "ttl-key", func() {})
assert.IsNil(err)
}
func TestRWKeyedLocker_LockAndUnlock(t *testing.T) {
t.Parallel()
assert := internal.NewAssert(t, "TestKeyedLocker_LockReleaseAfterTTL")
locker := NewRWKeyedLocker[string](500 * time.Millisecond)
var locked bool
err := locker.Lock(context.Background(), "key1", func() {
locked = true
})
assert.IsNil(err)
assert.Equal(true, locked)
}
func TestRWKeyedLocker_RLockParallel(t *testing.T) {
t.Parallel()
assert := internal.NewAssert(t, "TestKeyedLocker_LockReleaseAfterTTL")
locker := NewRWKeyedLocker[string](1 * time.Second)
var mu sync.Mutex
var count int
wg := sync.WaitGroup{}
for i := 0; i < 5; i++ {
wg.Add(1)
go func() {
defer wg.Done()
err := locker.RLock(context.Background(), "shared-key", func() {
time.Sleep(10 * time.Millisecond)
mu.Lock()
count++
mu.Unlock()
})
assert.IsNil(err)
}()
}
wg.Wait()
assert.Equal(5, count)
}
func TestRWKeyedLocker_LockTimeout(t *testing.T) {
t.Parallel()
assert := internal.NewAssert(t, "TestRWKeyedLocker_LockTimeout")
locker := NewRWKeyedLocker[string](1 * time.Second)
start := make(chan struct{})
go func() {
locker.Lock(context.Background(), "key-timeout", func() {
close(start)
time.Sleep(200 * time.Millisecond)
})
}()
<-start
ctx, cancel := context.WithTimeout(context.Background(), 50*time.Millisecond)
defer cancel()
err := locker.Lock(ctx, "key-timeout", func() {
t.Error("should not reach here")
})
assert.IsNotNil(err)
}

View File

@@ -1,51 +1,51 @@
-----BEGIN rsa private key-----
MIIJKgIBAAKCAgEAp8NkRHdZhqGYZ6tkWy3xjkm7PlRqi2yS+AgguuuWhw3DvhM4
BYwKt+r7Z0aEK9uIqpi2GwLv475kt+sQlPFbW/qVEM2RE9X4kZFhc56ZfomeUq/p
uUd6Yz4Ba7JkXsk1ErwtMbzf+eNPlCA9idBZOfccy6nXHQkGsQQO9o8AM9jdOfhI
8x/8CqJss2IZjQZxlzc37Rqs+OVD2S8S2Dr6QUPyD/2P1dCwm4imw7BX3BMcwVg/
LCQIGJIoMQ16T6prA/c6ShtnxHhc2ZyJ8ito7Bs9v7+9yGd6w6Img/I0QCbFwGbX
FSHLypFV2XauzFZbGQg7WM2WuKuU6NvRx1/EI/DMXUIXRNvAlgXYbJc1V+x/2718
MuuUaQjw39xderTkhU26wL2Rr/WfUaHQRq7FalaPohpw9P5awm2dZGRo+mw+KmzW
Mf/Yb2h11oVBeeHK1H/FoDBR8tq519QT2dSnU4teIXqfT8KR40PZQ8IzIISyGoXi
2uvlLIGrht6xweysyOAHHAaoDiX/ZeYDzhsdcRzo6fih+JPs/AwbeoNVuTBb0FZs
lEXmB+rWdlCABWyvKjXqFaBOw966muBoDpvlFhIXXT7cAhFN072yaDNa3judlY9y
iBAWQQsv2Qr05YJIfdS3nQ3bqIu3BIwOhv/tHCfUANM3DosKam4kxEIydWECAwEA
AQKCAgEAmcGBxyJfwf2e8fzqiIrOJiu7WgACenBzLrI8VTSQjIz4BuAUazkTpcbO
zbOadZvKYRh/ZrhFZsTcCJh/ZRLkOaOrNXBCdByaqcfFujL02T2GBqDFpJM3P3fX
0333ccwVQWuIPFqwKJXGHKuD2yhCbtbl5F9wEWNxZ5GhqSYc+GfdMkE1kuaQmKqO
18WkR1VNjFsGfeACAkgV7Bqxuc4sCN8eHru1NTEEfDg9J4Mas1As/aNEms8XQHXM
MlD49cTqOgM/wCXn7/CkoKlQ4MwaaLL64B/6746vvFeD11CHxPgELDfVDsAfyDN1
rE6THCJVcdbSlawvZGeVnenCFWnXF+G850MVMi3sPEIwsntItON/QdKYCR0MG4dQ
dVRRMhuA+k3YJ3cvjnYBMi/o9EyL3WX0rB5CZ/UGJdbJ9Fnlt3ow5z7Au73SR+m2
7G5xyoUz28qY0kpMPrYot3g2mKonzLc/bcWG7B1lt0L/rLsdnQ0bC0sV/JEnUO2D
3meI/bILIsBKx+Cb3kerGRxSoxY8mloEeksFok7lHKBdukZr7xpNUGWcsCCTbNNn
qCROwwX9vNQKmZGBjql5V8vvmucFEl9XivHZFNWbwETFK4lNBQ/KkjFFQvLTZqwZ
pkUI7xegzu3Lvg8dYohRqH2KUurO0x9Yz+2XH+v1NLgCLUAtcBUCggEBAM5+Za4e
q36C3MQ7QwgOEPDBmwFnn3dxJ8oa2MEdKMPSI3WGUsueNuw5+hrqy/u+pYMD5+/d
yQEHVgXopHsE8cygypw3StjgAxFSeVDwIOzcMSfcRkLqKPiarF63vl9Ng7ySZCxW
my6+Bc+IvexLZJ6pqclTgZS/aPlNRmGEgCy+k5SJFzsSD86qICjonff5ED/oLPnl
MNdo4nPcgpEaY3PMlsuLQFqf1SVgxIVzIbyqkn12Ed8NWEq7UQZIJsIJfljzZHUI
VqrijEsVzb65APs+YkINRae1Ni5UoHa6Rh/tBdFatiB13vr/gxzl//oejJuWQuhB
A6w90V3qroqXEwMCggEBAM/75r4zqDz/NzX5dVsGA/hll6SM2i6Y4F8c04CkprPM
2zCrkuDreXTNDK4+N/Hpi1G8fBgXqB6ygHiaCNcaHYiTpXvdnexwiHyqivoqWHq8
09f1OudFIkOA5VgApOoxZBnxlh3PCt2ULb+R37tKoCdg2KubzMcFxcMwiIL5yqaf
iBieQHvNKo2o6oQpezAwhudY/Ke5RoEpMl9Kav23lHbizyHZosuwqIlnyPzFlpIz
oAlmhLct9YCNwT5f9fBNRcvcG7J5lpTPGMjZIbSZWN5lvzngr/r6szHKAUjXS4R4
nmymUzN302ugrcrza5vOec1XCkvZnMbmvzh5T6zhdssCggEAPcPbADkWTPIxvNSJ
GVizwn/2sHXhYiXLpA1htmnVbrVle9rg2x0JCqHQ6MpAl52P/l9luf7aB9+84GmV
AWMaPH3//LghQDvJTx4fQZGCF7dJUMX3kj5eYPZLBs3pOLKD7BzRr91774BRVqFt
RcfLYhYXviunP+n9KUzu925dtISQukZDV5zwc325vuLNlYW/UY5OHbVrZZNu4P8d
Yu/Evbd9h4awiiry44pNXilw9vECflqZv+FK/peHBd0BEtsqGss6yjLjUZwQIMl1
0E2gOIaRd3Zm1mJCwZr4oGrZXOVV7yg2AAGh3+wbuMInThZjMorAmp4Pzi5zeKcg
7D2CJwKCAQEAvW2m6VFPR88DUCuWkYLXFuQgy4RmK86dfMNad25/Tn+km52JN0YA
5zri88hDWBfoBlfvhln1i4/0puNUbeWhRIWFUV21umV8Sl9iGRt8Xr2hDK6UKZOz
81twh0h+67Z0f5Sjrx6lvM57JGIOLh135KW2cgaC6jn9txt7Gh+8TSo74IAyJw/k
VAWnIxxM0MVB+W/5HiIHbxhAgr3a2J7dn7JQCXqZZX/O4OcgDelAjIRsnGM9OUGo
up9hhBgOfgFDMruUlmdbmMlOv4/TvRN194kgM+zLG4I+t3hO1zMP2uWpFTgfy161
tu8vmws91TingzhlblQTEK2VODB1OvZXJQKCAQEAqAN1EiRfN++Kpz3RZUMqz2GT
yoy614rFvut1tCTicLzBmHZ+xI5HNknXqTQZVgxJk5iB4Vbq50VxcunVQMn7bcm8
AgGP/tH27prTX9KwbYuX3uaI58/F7Ir2iP3fF3eb0ejqWrnxCnC+hm9Up9Ivrfvy
8r7lJjNdmVquhLf/MUCQnbZcFdeh3B7orWmf9CdfJcCYhGbhBSk2xYuwRXqXb/A3
Uelrjur25182esKPoInpM+STfGIF4WgCR8lo0npjGSz0oiBiyiJTpcS/AqVKYgSw
5kVP4jUprOqKYVO7bL74qiYzsrujJh2+XfDUApsyhvn63fYRUeygTQllDysyYw==
MIIJKQIBAAKCAgEAtKEukXkdl8ggX5I5Us0MBcsrFzbTgwOn8u72IhJzruvsHcjl
GGF5L1eXh3C0zbcoWq+Bxf8NnXvSXASZlZHTiSVLFerxFw2LwTNcyfOz/Y/wIbtr
Yu4ZdttLKdJS607UvnZ57O7SQni5rxYAP5j+veKHEJgAO7Sn0lEEWwkeudsfVrQm
Qx6JHYtCtx78p2GHXPj9wXWHxWHKMkczKtmBFDLlYocPhQfWzx87sYKjGmTMGzVd
dsrllAqVozlxXTQndOdozx8BWHSBiYWguOvWiYjvryrIf0B7wn4wt3o2Yre3w4Bk
wrTUE73JWG2q1NLuPoT27m7WcdVFGGcKAygAdMMSwS530v8tnAw876/w1SfUmtkI
cqknysMjsvxHx1fbEfkl3b0yBnzcUrT1qxzmNEqJSQXBtFYFyGgyeL8J1ul/T5hP
SNz0we5TO/xRkX5Y1HNJ+a2r8EnE1PU8zgQGBxYSIub4ajSUccyn5yHQ/k2Flsol
Ziu/UzfnzMgdg55svgf5x1ntKYBoVHuYe84Ab3WRAeVlUi0OphvFlBzfIQvrG2sR
biPZRZZ0E8kBH94zPmrFmyX+7TwlSqVHtNFc+Ynz+hrHcd4z6N7fIb+D1gSYuukk
cRqr4ANO0ugZk0YYKaWDf86H7OZOXna8ObIAi+3gx+Wcp1lm7Ffk9sn+dGkCAwEA
AQKCAgAdp/kFWWVKbkkiZ9eRiKjYqqrAfPftIsSIVkODBJSJu6JgoYM7pYVICJGQ
YyjMPa4adYZRA7cwjAvVn8u3iuG4Oq9BQfmjV04CwnQRlDmQ//jlEOhorb7wjMCi
dS24BJFZVApgpDVRRJD39hzEVFI+ytpyFwKyys5i8XpNmAm7agaTLbC6hGDuwNaL
SkMhGBopYZgIE0vfVFbmOlpkRqGyt0iCDLq3lLnn97DNTC2LP9FjBjf6MQXQcIxw
6BV8v/tabkP+/ZAy/a3m7lGdCtuGaT0w/U0911B6dk15Uk9rlc5OAt6IOTg4pYhR
RHAv0RHcoegI7Zm3xtQ5VXGHYyHvwwXA/e5pAqOOO/GZ+85OcilphnSLGoP71rqJ
h22GXTj4uCWiiPA8yLwAdtLKhN/KygzH3gkrsD5bH5R9ZmnBiIk192X2v1flND8M
ikBx6ADCphQnJ2zhCv6teesro9A0GIHVI0NjpZayvKjHMHpMpHpN8ZLMTmbafIH3
tq3g9H17ncWal2qDKco7J+6VgJx1NRGlQY/52J25BX2UzB3rrFW21w01TCtpKH3d
o4sPl5IQS5RsrlPEIVAOo7tYHsB5pgfWUzMN3u9hONnES4NHwgDgYJlRl9B/unPh
/kbz/lEJq3Cw2F3zhYb/8neM49lgmWq+hpUshejiYABaXpMGoQKCAQEA4audTQEt
cH89v7NGWWUplfIgb2ndXUW8IrJQbuLbFkCdlXbSfsOA6GsA8bMNI9IUqAfCpvNg
qmCdJ7b8kfHDIilW2SPDA9wW3Fd30ofJEIkfRet87DeT541E27hUHTZ5GSiKvzqn
j+IXu2mpcASgeOJnIp6KmyTOIwj9DJZPexnDRYwnbjpLUdWCCkGcXbIMggrvjaTj
QEWW8RWt+GPyhDACh1KFa2RSON5p3HasPKZJmJz6Zgk9rJ7xI3MSdnGjGMl4PviB
LyrXjf5URqrDmpUf+Xtl3ysNiGDzRBn05SjpmsS2yOjQlK5CiK/o+LTIsNUtdzay
5PLoxR74b6uTtwKCAQEAzOfmshWZ7AbaaLYPCdV3aDRsti4O8U400cc4QFTMBjwZ
AI8A/++fcTZJgtmFPMhg4C9BNj1w8sgcQ7otTv5iJFgZ2bktBp5dgWsX4SbyrC16
xapOa1PNcgV+YmZ2HnzC4lCfE42hsTP2QI2E2Yr1vyeZ9YvbWJHvBJYA0s5lwkpp
dO7abJ/F6DYzpvjomAobutTQmd0XvNTSLuhfudRDadC4RtBQ/bjF+Rx3EjXeWJzy
l83VBsDnlqhrsGpV+yq1pRePfY9XzDjWfz+oOzCMSvl1+eXmFxl7PJqXc9JD2CeW
6xfzQEXxB7czHaTJ4tSB4yC7U9qOXLkZQBpyGg143wKCAQA9raH4gfHhZWWDF4SK
ulN7YAntaYnPDFg3Q3UoWWh31IE9cJRngReiblx7suxMdgafRj+1UZ+B8ZYCXMj7
OpCSranG/zc1vtmgr2dYazRRCKk7evlRtn7+MmY3h1G2CkVe0u3ZBjb15F2II4Dj
1N/nKjn2BE7tyElu2e4PmqVuh8QPJhdA0T30x94a34PVN+yjPknq9L4Huv1eNwat
dOO7rUODqNI+X9T5JhDY6LZ6fRhwVbc6XBw3KdnOTo0lQjnJdIcg7tqgAZ2YeYKf
Ldz4SvnKPifBrwqr05OpcU61s1DltA4hK0CW4mnc4fdSwlZ3vkwG4TRTzvA/sA9G
tiZRAoIBAQCuVmCiBF8BwpLxpHUHGOiPcItONcHg7XljQu1JTtyIMXnUT9e56lbu
LBI/knMaVCKYm5wQWhZPepMRzMXf/+/gnFTiOftlNji4dDXNCyZN+CQNKemux451
BNeTQToelmf5xj6SlF6ONne+VKpDrUeJbFhB4sytfvyuGjJ5KcLKnCU9qDuPUCFC
gVtRJVZAhdkyDP+u6b3Ym/p4jp1jroXs8fjXx0YhmaRXXzCv/cU//8kn/6jQJjDk
rkdxwgeFu8Dwxir/2YYJ7BIUEkVAlv3GjJkkFca+wJ9p4N4bXTr8HjL5s1bzyI5a
0jRbdGmQ5N3eMWsw3TNjENm7AMU0BWJhAoIBAQCdW0LoiFahv/zrqM/gBE88C1fS
w4/BrSl/bhLGZQqSbcr9SxM/2Y5xyE4RS52glSDfx4zlseI93SqQDKLfGP+bjI9d
HG07T3SELvs816SFRuGC+ktw3mFjFaGn4KlLtOsGcPRMZofsm/XdUNpWj5Ic4d5/
4YWwBQn9VaODDvrhCcv+biA0vRBPlH67gubCy8RMb7J2GspMU9ekIjQxHJkVLY/a
k9gBVdAtzeQY9VHLsc5/zTXReXUCyvrUw+h2hVtGx3CgncKH4WWELBWHXZqQJ6kJ
sZhZ+KH9b6XGnSY29wPX1GKCTIXXIfN5xzFo65yw2b8WE0u0idvEsv1oKqtC
-----END rsa private key-----

View File

@@ -1,14 +1,14 @@
-----BEGIN rsa public key-----
MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEAp8NkRHdZhqGYZ6tkWy3x
jkm7PlRqi2yS+AgguuuWhw3DvhM4BYwKt+r7Z0aEK9uIqpi2GwLv475kt+sQlPFb
W/qVEM2RE9X4kZFhc56ZfomeUq/puUd6Yz4Ba7JkXsk1ErwtMbzf+eNPlCA9idBZ
Ofccy6nXHQkGsQQO9o8AM9jdOfhI8x/8CqJss2IZjQZxlzc37Rqs+OVD2S8S2Dr6
QUPyD/2P1dCwm4imw7BX3BMcwVg/LCQIGJIoMQ16T6prA/c6ShtnxHhc2ZyJ8ito
7Bs9v7+9yGd6w6Img/I0QCbFwGbXFSHLypFV2XauzFZbGQg7WM2WuKuU6NvRx1/E
I/DMXUIXRNvAlgXYbJc1V+x/2718MuuUaQjw39xderTkhU26wL2Rr/WfUaHQRq7F
alaPohpw9P5awm2dZGRo+mw+KmzWMf/Yb2h11oVBeeHK1H/FoDBR8tq519QT2dSn
U4teIXqfT8KR40PZQ8IzIISyGoXi2uvlLIGrht6xweysyOAHHAaoDiX/ZeYDzhsd
cRzo6fih+JPs/AwbeoNVuTBb0FZslEXmB+rWdlCABWyvKjXqFaBOw966muBoDpvl
FhIXXT7cAhFN072yaDNa3judlY9yiBAWQQsv2Qr05YJIfdS3nQ3bqIu3BIwOhv/t
HCfUANM3DosKam4kxEIydWECAwEAAQ==
MIICIjANBgkqhkiG9w0BAQEFAAOCAg8AMIICCgKCAgEAtKEukXkdl8ggX5I5Us0M
BcsrFzbTgwOn8u72IhJzruvsHcjlGGF5L1eXh3C0zbcoWq+Bxf8NnXvSXASZlZHT
iSVLFerxFw2LwTNcyfOz/Y/wIbtrYu4ZdttLKdJS607UvnZ57O7SQni5rxYAP5j+
veKHEJgAO7Sn0lEEWwkeudsfVrQmQx6JHYtCtx78p2GHXPj9wXWHxWHKMkczKtmB
FDLlYocPhQfWzx87sYKjGmTMGzVddsrllAqVozlxXTQndOdozx8BWHSBiYWguOvW
iYjvryrIf0B7wn4wt3o2Yre3w4BkwrTUE73JWG2q1NLuPoT27m7WcdVFGGcKAygA
dMMSwS530v8tnAw876/w1SfUmtkIcqknysMjsvxHx1fbEfkl3b0yBnzcUrT1qxzm
NEqJSQXBtFYFyGgyeL8J1ul/T5hPSNz0we5TO/xRkX5Y1HNJ+a2r8EnE1PU8zgQG
BxYSIub4ajSUccyn5yHQ/k2FlsolZiu/UzfnzMgdg55svgf5x1ntKYBoVHuYe84A
b3WRAeVlUi0OphvFlBzfIQvrG2sRbiPZRZZ0E8kBH94zPmrFmyX+7TwlSqVHtNFc
+Ynz+hrHcd4z6N7fIb+D1gSYuukkcRqr4ANO0ugZk0YYKaWDf86H7OZOXna8ObIA
i+3gx+Wcp1lm7Ffk9sn+dGkCAwEAAQ==
-----END rsa public key-----