Skip to content

Commit f02eadb

Browse files
michaelstaibmichaelstaib
authored
Optimize locking on AsyncAutoResetEvent (#9430)
1 parent b70e7d6 commit f02eadb

2 files changed

Lines changed: 132 additions & 106 deletions

File tree

‎src/HotChocolate/Core/src/Types/Fetching/AsyncAutoResetEvent.cs‎

Lines changed: 66 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -5,69 +5,90 @@ namespace HotChocolate.Fetching;
55

66
internal sealed class AsyncAutoResetEvent : INotifyCompletion
77
{
8-
#if NET9_0_OR_GREATER
9-
private readonly Lock _sync = new();
10-
#else
11-
private readonly object _sync = new();
12-
#endif
8+
private const int Idle = 0;
9+
private const int Signaled = 1;
10+
private const int Waiting = 2;
11+
12+
private int _state;
1313
private Action? _continuation;
14-
private bool _isSignaled;
1514

16-
public bool IsSignaled => Volatile.Read(ref _isSignaled);
15+
public bool IsSignaled => Volatile.Read(ref _state) == Signaled;
1716

18-
public bool IsCompleted => false;
17+
public bool IsCompleted => Volatile.Read(ref _state) == Signaled;
1918

20-
public void GetResult() { }
19+
public void GetResult()
20+
{
21+
// Consume the signal when completing synchronously (via IsCompleted == true).
22+
// CAS failure is benign (e.g. TryResetToIdle cleared it first).
23+
Interlocked.CompareExchange(ref _state, Idle, Signaled);
24+
}
2125

2226
public void OnCompleted(Action continuation)
2327
{
24-
bool wasSignaled;
28+
Debug.Assert(_continuation is null, "There should only be one awaiter.");
29+
_continuation = continuation;
2530

26-
lock (_sync)
31+
while (true)
2732
{
28-
wasSignaled = _isSignaled;
29-
30-
if (wasSignaled)
31-
{
32-
// consume the signal
33-
_isSignaled = false;
34-
}
35-
else
33+
switch (Volatile.Read(ref _state))
3634
{
37-
Debug.Assert(_continuation is null, "There should only be one awaiter.");
38-
_continuation = continuation;
39-
}
40-
}
35+
case Idle:
36+
// Register waiter: IDLE -> WAITING
37+
if (Interlocked.CompareExchange(ref _state, Waiting, Idle) == Idle)
38+
{
39+
return;
40+
}
41+
break; // CAS failed, retry
4142

42-
if (wasSignaled)
43-
{
44-
ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: true);
43+
case Signaled:
44+
// Consume signal immediately: SIGNALED -> IDLE
45+
if (Interlocked.CompareExchange(ref _state, Idle, Signaled) == Signaled)
46+
{
47+
_continuation = null;
48+
ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: true);
49+
return;
50+
}
51+
break; // CAS failed, retry
52+
53+
default:
54+
Debug.Fail("OnCompleted called while already waiting.");
55+
return;
56+
}
4557
}
4658
}
4759

4860
public void Set()
4961
{
50-
Action? continuation = null;
51-
52-
lock (_sync)
62+
while (true)
5363
{
54-
if (_continuation is not null)
64+
switch (Volatile.Read(ref _state))
5565
{
56-
// someone is waiting - release them immediately
57-
// we don't set _isSignaled since we're consuming it immediately
58-
continuation = _continuation;
59-
_continuation = null;
60-
}
61-
else
62-
{
63-
// since no one waiting we are storing the signal for the next awaiter
64-
_isSignaled = true;
65-
}
66-
}
66+
case Idle:
67+
// Store signal: IDLE -> SIGNALED
68+
if (Interlocked.CompareExchange(ref _state, Signaled, Idle) == Idle)
69+
{
70+
return;
71+
}
72+
break; // CAS failed, retry
6773

68-
if (continuation is not null)
69-
{
70-
ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: true);
74+
case Waiting:
75+
// Wake waiter: WAITING -> IDLE
76+
if (Interlocked.CompareExchange(ref _state, Idle, Waiting) == Waiting)
77+
{
78+
var c = _continuation!;
79+
_continuation = null;
80+
ThreadPool.QueueUserWorkItem(static c => c(), c, preferLocal: true);
81+
return;
82+
}
83+
break; // CAS failed, retry
84+
85+
case Signaled:
86+
// Already signaled, nothing to do
87+
return;
88+
89+
default:
90+
return;
91+
}
7192
}
7293
}
7394

@@ -77,15 +98,7 @@ public void Set()
7798
/// </summary>
7899
public bool TryResetToIdle()
79100
{
80-
lock (_sync)
81-
{
82-
if (_continuation is null && _isSignaled)
83-
{
84-
_isSignaled = false;
85-
return true;
86-
}
87-
return false;
88-
}
101+
return Interlocked.CompareExchange(ref _state, Idle, Signaled) == Signaled;
89102
}
90103

91104
public AsyncAutoResetEvent GetAwaiter() => this;

‎src/HotChocolate/Fusion/src/Fusion.Execution/Execution/AsyncAutoResetEvent.cs‎

Lines changed: 66 additions & 53 deletions
Original file line numberDiff line numberDiff line change
@@ -5,69 +5,90 @@ namespace HotChocolate.Fusion.Execution;
55

66
internal sealed class AsyncAutoResetEvent : INotifyCompletion
77
{
8-
#if NET9_0_OR_GREATER
9-
private readonly Lock _sync = new();
10-
#else
11-
private readonly object _sync = new();
12-
#endif
8+
private const int Idle = 0;
9+
private const int Signaled = 1;
10+
private const int Waiting = 2;
11+
12+
private int _state;
1313
private Action? _continuation;
14-
private bool _isSignaled;
1514

16-
public bool IsSignaled => Volatile.Read(ref _isSignaled);
15+
public bool IsSignaled => Volatile.Read(ref _state) == Signaled;
1716

18-
public bool IsCompleted => false;
17+
public bool IsCompleted => Volatile.Read(ref _state) == Signaled;
1918

20-
public void GetResult() { }
19+
public void GetResult()
20+
{
21+
// Consume the signal when completing synchronously (via IsCompleted == true).
22+
// CAS failure is benign (e.g. TryResetToIdle cleared it first).
23+
Interlocked.CompareExchange(ref _state, Idle, Signaled);
24+
}
2125

2226
public void OnCompleted(Action continuation)
2327
{
24-
bool wasSignaled;
28+
Debug.Assert(_continuation is null, "There should only be one awaiter.");
29+
_continuation = continuation;
2530

26-
lock (_sync)
31+
while (true)
2732
{
28-
wasSignaled = _isSignaled;
29-
30-
if (wasSignaled)
31-
{
32-
// consume the signal
33-
_isSignaled = false;
34-
}
35-
else
33+
switch (Volatile.Read(ref _state))
3634
{
37-
Debug.Assert(_continuation is null, "There should only be one awaiter.");
38-
_continuation = continuation;
39-
}
40-
}
35+
case Idle:
36+
// Register waiter: IDLE -> WAITING
37+
if (Interlocked.CompareExchange(ref _state, Waiting, Idle) == Idle)
38+
{
39+
return;
40+
}
41+
break; // CAS failed, retry
4142

42-
if (wasSignaled)
43-
{
44-
ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: true);
43+
case Signaled:
44+
// Consume signal immediately: SIGNALED -> IDLE
45+
if (Interlocked.CompareExchange(ref _state, Idle, Signaled) == Signaled)
46+
{
47+
_continuation = null;
48+
ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: true);
49+
return;
50+
}
51+
break; // CAS failed, retry
52+
53+
default:
54+
Debug.Fail("OnCompleted called while already waiting.");
55+
return;
56+
}
4557
}
4658
}
4759

4860
public void Set()
4961
{
50-
Action? continuation = null;
51-
52-
lock (_sync)
62+
while (true)
5363
{
54-
if (_continuation is not null)
64+
switch (Volatile.Read(ref _state))
5565
{
56-
// someone is waiting - release them immediately
57-
// we don't set _isSignaled since we're consuming it immediately
58-
continuation = _continuation;
59-
_continuation = null;
60-
}
61-
else
62-
{
63-
// since no one waiting we are storing the signal for the next awaiter
64-
_isSignaled = true;
65-
}
66-
}
66+
case Idle:
67+
// Store signal: IDLE -> SIGNALED
68+
if (Interlocked.CompareExchange(ref _state, Signaled, Idle) == Idle)
69+
{
70+
return;
71+
}
72+
break; // CAS failed, retry
6773

68-
if (continuation is not null)
69-
{
70-
ThreadPool.QueueUserWorkItem(static c => c(), continuation, preferLocal: true);
74+
case Waiting:
75+
// Wake waiter: WAITING -> IDLE
76+
if (Interlocked.CompareExchange(ref _state, Idle, Waiting) == Waiting)
77+
{
78+
var c = _continuation!;
79+
_continuation = null;
80+
ThreadPool.QueueUserWorkItem(static c => c(), c, preferLocal: true);
81+
return;
82+
}
83+
break; // CAS failed, retry
84+
85+
case Signaled:
86+
// Already signaled, nothing to do
87+
return;
88+
89+
default:
90+
return;
91+
}
7192
}
7293
}
7394

@@ -77,15 +98,7 @@ public void Set()
7798
/// </summary>
7899
public bool TryResetToIdle()
79100
{
80-
lock (_sync)
81-
{
82-
if (_continuation is null && _isSignaled)
83-
{
84-
_isSignaled = false;
85-
return true;
86-
}
87-
return false;
88-
}
101+
return Interlocked.CompareExchange(ref _state, Idle, Signaled) == Signaled;
89102
}
90103

91104
public AsyncAutoResetEvent GetAwaiter() => this;

0 commit comments

Comments
 (0)