Skip to content

Commit 6ab5d4e

Browse files
arenchenarenchen
committed
improve
1 parent 25dca2a commit 6ab5d4e

6 files changed

Lines changed: 259 additions & 138 deletions

File tree

‎KeyCountLocker.cs‎

Lines changed: 0 additions & 76 deletions
This file was deleted.

‎KeyLocker.cs‎

Lines changed: 223 additions & 50 deletions
Original file line numberDiff line numberDiff line change
@@ -1,95 +1,268 @@
11
using System.Collections.Concurrent;
2-
using System.Diagnostics;
32
using Microsoft.Extensions.Logging;
43

54
namespace Netcorext.Extensions.Threading;
65

7-
public class KeyLocker
6+
public class KeyLocker : IDisposable
87
{
9-
public const long DEFAULT_DEAD_LOCK_TIMES = 100;
8+
private bool _disposed;
9+
private readonly int _maxConcurrent;
10+
private readonly TimeSpan? _timeout;
11+
private readonly bool _throwTimeoutException;
12+
private readonly TimeSpan? _expired;
13+
private readonly ILogger _logger;
14+
private readonly ConcurrentDictionary<string, KeyState> _locks = new();
15+
private readonly ConcurrentDictionary<string, Timer> _timers = new();
1016

11-
private readonly long _deadLockTimes = DEFAULT_DEAD_LOCK_TIMES;
12-
private readonly ILogger? _logger;
13-
private static readonly ConcurrentDictionary<string, KeyLockerState<bool>> Lockers = new();
14-
15-
public KeyLocker()
16-
{ }
17-
18-
public KeyLocker(long deadLockTimes = DEFAULT_DEAD_LOCK_TIMES, ILogger? logger = null)
17+
public KeyLocker(ILogger logger, TimeSpan? timeout = null, bool throwTimeoutException = false, TimeSpan? expired = null, int maxConcurrent = 1)
1918
{
20-
_deadLockTimes = deadLockTimes;
19+
_maxConcurrent = maxConcurrent;
20+
_timeout = timeout;
21+
_throwTimeoutException = throwTimeoutException;
22+
_expired = expired ?? TimeSpan.FromMilliseconds(10 * 60 * 1000);
2123
_logger = logger;
2224
}
2325

24-
public async Task WaitAsync(string key, CancellationToken cancellationToken = default)
26+
public void Wait(string key)
2527
{
26-
await WaitAsync(key, null, false, cancellationToken);
28+
var keyState = _locks.AddOrUpdate(key, CreateLockItem, (k, state) =>
29+
{
30+
lock (state)
31+
{
32+
if (!state.ReleaseAll && !state.Cancellation.IsCancellationRequested)
33+
state.IncrementConcurrent();
34+
35+
state.LastWaitingTime = DateTimeOffset.UtcNow;
36+
}
37+
38+
return state;
39+
});
40+
41+
try
42+
{
43+
if (keyState.ReleaseAll)
44+
return;
45+
46+
if (_expired.HasValue && _timers.TryAdd(key, new Timer(HandleExpired, key, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan)))
47+
_timers[key].Change(_expired.Value, _expired.Value);
48+
49+
if (_timeout.HasValue)
50+
{
51+
if (keyState.Semaphore.Wait(_timeout.Value, keyState.Cancellation.Token))
52+
return;
53+
54+
HandleTimeout(keyState);
55+
}
56+
else
57+
{
58+
keyState.Semaphore.Wait(keyState.Cancellation.Token);
59+
}
60+
}
61+
catch (OperationCanceledException) { }
2762
}
2863

29-
public async Task WaitAsync(string key, TimeSpan? timeout = null, CancellationToken cancellationToken = default)
64+
public async Task WaitAsync(string key)
3065
{
31-
await WaitAsync(key, timeout, false, cancellationToken);
66+
var keyState = _locks.AddOrUpdate(key, CreateLockItem, (k, state) =>
67+
{
68+
lock (state)
69+
{
70+
if (!state.ReleaseAll && !state.Cancellation.IsCancellationRequested)
71+
state.IncrementConcurrent();
72+
73+
state.LastWaitingTime = DateTimeOffset.UtcNow;
74+
}
75+
76+
return state;
77+
});
78+
79+
try
80+
{
81+
if (keyState.ReleaseAll)
82+
return;
83+
84+
if (_expired.HasValue && _timers.TryAdd(key, new Timer(HandleExpired, key, Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan)))
85+
_timers[key].Change(_expired.Value, _expired.Value);
86+
87+
if (_timeout.HasValue)
88+
{
89+
if (await keyState.Semaphore.WaitAsync(_timeout.Value, keyState.Cancellation.Token))
90+
return;
91+
92+
HandleTimeout(keyState);
93+
}
94+
else
95+
{
96+
await keyState.Semaphore.WaitAsync(keyState.Cancellation.Token);
97+
}
98+
}
99+
catch (OperationCanceledException) { }
32100
}
33101

34-
public async Task WaitAsync(string key, TimeSpan? timeout = null, bool releaseAll = false, CancellationToken cancellationToken = default)
102+
public int Release(string key)
35103
{
36-
var stopwatch = new Stopwatch();
104+
if (!_locks.TryGetValue(key, out var keyState))
105+
return 0;
106+
107+
lock (keyState)
108+
{
109+
if (keyState.ReleaseAll || keyState.Cancellation.IsCancellationRequested)
110+
return 0;
37111

38-
stopwatch.Start();
112+
try
113+
{
114+
keyState.LastWaitingTime = DateTimeOffset.UtcNow;
115+
keyState.Semaphore.Release();
116+
keyState.DecrementConcurrent();
39117

40-
while (!Lockers.TryAdd(key, new KeyLockerState<bool>
41-
{
42-
State = true,
43-
Expires = timeout.HasValue ? DateTimeOffset.UtcNow.Add(timeout.Value) : null
44-
}) && Lockers.TryGetValue(key, out var locker) && locker.State)
118+
return 1;
119+
}
120+
catch (ObjectDisposedException) { }
121+
catch (ArgumentOutOfRangeException) { }
122+
catch (SemaphoreFullException) { }
123+
}
124+
125+
return 0;
126+
}
127+
128+
public int ReleaseAll(string key)
129+
{
130+
if (!_locks.TryGetValue(key, out var keyState))
131+
return 0;
132+
133+
lock (keyState)
45134
{
46-
if (locker.IsExpired)
135+
if (keyState.ReleaseAll || keyState.Cancellation.IsCancellationRequested)
136+
return 0;
137+
138+
keyState.ReleaseAll = true;
139+
140+
var releaseCount = 0;
141+
142+
try
47143
{
48-
if (releaseAll)
144+
while (true)
49145
{
50-
ReleaseAll(key);
51-
}
52-
else
53-
{
54-
Release(key);
146+
try
147+
{
148+
keyState.LastWaitingTime = DateTimeOffset.UtcNow;
149+
keyState.Semaphore.Release();
150+
keyState.DecrementConcurrent();
151+
releaseCount++;
152+
}
153+
catch (SemaphoreFullException)
154+
{
155+
keyState.Cancellation.Cancel(true);
156+
157+
break;
158+
}
55159
}
56160
}
161+
catch (ObjectDisposedException) { }
162+
catch (ArgumentOutOfRangeException) { }
163+
164+
return releaseCount;
165+
}
166+
}
167+
168+
public void Reset(string key)
169+
{
170+
if (!_locks.TryGetValue(key, out var keyState))
171+
return;
57172

58-
if (_logger != null && stopwatch.ElapsedMilliseconds >= _deadLockTimes && stopwatch.ElapsedMilliseconds % _deadLockTimes == 0)
59-
_logger.LogWarning("'{Key}' locked for too long, elapsed: {StopwatchElapsed}", key, stopwatch.Elapsed);
173+
lock (keyState)
174+
{
175+
keyState.LastWaitingTime = DateTimeOffset.UtcNow;
176+
keyState.Cancellation = new CancellationTokenSource();
177+
keyState.ReleaseAll = false;
178+
}
60179

61-
await Task.Delay(1, cancellationToken);
180+
if (!_timers.TryRemove(key, out var timer))
181+
return;
182+
183+
lock (timer)
184+
{
185+
timer.Change(Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
186+
timer.Dispose();
62187
}
188+
}
63189

64-
stopwatch.Stop();
190+
public int GetWaitingCount(string key)
191+
{
192+
return _locks.TryGetValue(key, out var keyState) ? keyState.WaitingConcurrent : 0;
65193
}
66194

67-
public bool Release(string key)
195+
private void HandleTimeout(KeyState keyState)
68196
{
69-
return Lockers.TryRemove(key, out _) || !Lockers.ContainsKey(key);
197+
_logger.LogWarning("Lock on key '{Key}' timed out", keyState.Key);
198+
199+
Release(keyState.Key);
200+
201+
if (_throwTimeoutException)
202+
throw new TimeoutException($"Lock on key '{keyState.Key}' timed out.");
70203
}
71204

72-
public bool ReleaseAll(string key)
205+
private void HandleExpired(object? state)
73206
{
74-
var newLocker = new KeyLockerState<bool>
75-
{
76-
State = false
77-
};
207+
if (state is not string key || !_locks.TryGetValue(key, out var keyState))
208+
return;
78209

79-
if (!Lockers.TryGetValue(key, out var oldLocker))
80-
return Lockers.TryAdd(key, newLocker);
210+
var elapsed = DateTimeOffset.UtcNow.Subtract(keyState.LastWaitingTime);
81211

82-
return Lockers.TryUpdate(key, newLocker, oldLocker);
212+
if (!_expired.HasValue || !(elapsed >= _expired))
213+
return;
214+
215+
lock (keyState)
216+
{
217+
if (!_locks.TryRemove(keyState.Key, out _))
218+
return;
219+
220+
_logger.LogWarning("Key '{Key}' expired({Elapsed}), has been removed", keyState.Key, elapsed);
221+
222+
if (_timers.TryRemove(keyState.Key, out var timer))
223+
{
224+
lock (timer)
225+
{
226+
timer.Change(Timeout.InfiniteTimeSpan, Timeout.InfiniteTimeSpan);
227+
timer.Dispose();
228+
}
229+
}
230+
231+
keyState.Dispose();
232+
}
233+
}
234+
235+
private KeyState CreateLockItem(string key)
236+
{
237+
var keyState = new KeyState
238+
{
239+
Key = key,
240+
Semaphore = new SemaphoreSlim(_maxConcurrent, _maxConcurrent),
241+
Cancellation = new CancellationTokenSource(),
242+
LastWaitingTime = DateTimeOffset.UtcNow,
243+
ReleaseAll = false
244+
};
245+
246+
return keyState;
83247
}
84248

85-
public void Prune(params string[] keys)
249+
public void Dispose()
86250
{
87-
var keysToRemove = keys.Length == 0 ? Lockers.Keys : keys;
251+
if (_disposed)
252+
return;
253+
254+
_disposed = true;
255+
256+
foreach (var key in _locks.Keys)
257+
{
258+
if (_locks.TryRemove(key, out var keyState))
259+
keyState.Dispose();
260+
}
88261

89-
foreach (var key in keysToRemove)
262+
foreach (var key in _timers.Keys)
90263
{
91-
if (Lockers.TryGetValue(key, out var locker) && locker.IsExpired)
92-
Lockers.TryRemove(key, out _);
264+
if (_timers.TryRemove(key, out var timer))
265+
timer.Dispose();
93266
}
94267
}
95268
}

0 commit comments

Comments
 (0)