Skip to content

Commit 95ae533

Browse files
committed
Use ManualResetValueTaskSourceCore instead of copying source code.
1 parent 44f986f commit 95ae533

File tree

1 file changed

+18
-119
lines changed

1 file changed

+18
-119
lines changed
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,4 @@
1-
using JetBrains.Annotations;
2-
using System;
3-
using System.Diagnostics;
4-
using System.Runtime.ExceptionServices;
5-
using System.Threading;
1+
using System;
62
using System.Threading.Tasks;
73
using System.Threading.Tasks.Sources;
84

@@ -13,140 +9,43 @@ namespace BenchmarkDotNet.Helpers
139
/// </summary>
1410
public class AutoResetValueTaskSource<TResult> : IValueTaskSource<TResult>, IValueTaskSource
1511
{
16-
[CanBeNull] private Action<object> _continuation;
17-
[CanBeNull] private object _continuationState;
18-
private bool _completed;
19-
[CanBeNull] private TResult _result;
20-
[CanBeNull] private ExceptionDispatchInfo _error;
21-
private short _version;
22-
23-
private void Reset()
24-
{
25-
// Reset/update state for the next use/await of this instance.
26-
_version++;
27-
_completed = false;
28-
_result = default;
29-
_error = null;
30-
_continuation = null;
31-
_continuationState = null;
32-
}
12+
private ManualResetValueTaskSourceCore<TResult> _sourceCore;
3313

3414
/// <summary>Completes with a successful result.</summary>
3515
/// <param name="result">The result.</param>
36-
public void SetResult(TResult result)
37-
{
38-
_result = result;
39-
SignalCompletion();
40-
}
16+
public void SetResult(TResult result) => _sourceCore.SetResult(result);
4117

4218
/// <summary>Completes with an error.</summary>
4319
/// <param name="error">The exception.</param>
44-
public void SetException(Exception error)
45-
{
46-
_error = ExceptionDispatchInfo.Capture(error);
47-
SignalCompletion();
48-
}
20+
public void SetException(Exception error) => _sourceCore.SetException(error);
4921

5022
/// <summary>Gets the operation version.</summary>
51-
public short Version => _version;
52-
53-
/// <summary>Gets the status of the operation.</summary>
54-
/// <param name="token">Opaque value that was provided to the <see cref="ValueTask"/>'s constructor.</param>
55-
public ValueTaskSourceStatus GetStatus(short token)
56-
{
57-
ValidateToken(token);
58-
return
59-
_continuation == null || !_completed ? ValueTaskSourceStatus.Pending :
60-
_error == null ? ValueTaskSourceStatus.Succeeded :
61-
_error.SourceException is OperationCanceledException ? ValueTaskSourceStatus.Canceled :
62-
ValueTaskSourceStatus.Faulted;
63-
}
23+
public short Version => _sourceCore.Version;
6424

65-
/// <summary>Gets the result of the operation.</summary>
66-
/// <param name="token">Opaque value that was provided to the <see cref="ValueTask"/>'s constructor.</param>
67-
public TResult GetResult(short token)
25+
private TResult GetResult(short token)
6826
{
69-
ValidateToken(token);
70-
if (!_completed)
27+
// We don't want to reset this if the token is invalid.
28+
if (token != Version)
7129
{
7230
throw new InvalidOperationException();
7331
}
74-
75-
var result = _result;
76-
var error = _error;
77-
Reset();
78-
error?.Throw();
79-
return result;
80-
}
81-
82-
public void OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags)
83-
{
84-
if (continuation == null)
85-
{
86-
throw new ArgumentNullException(nameof(continuation));
87-
}
88-
ValidateToken(token);
89-
90-
// We need to set the continuation state before we swap in the delegate, so that
91-
// if there's a race between this and SetResult/Exception and SetResult/Exception
92-
// sees the _continuation as non-null, it'll be able to invoke it with the state
93-
// stored here. However, this also means that if this is used incorrectly (e.g.
94-
// awaited twice concurrently), _continuationState might get erroneously overwritten.
95-
// To minimize the chances of that, we check preemptively whether _continuation
96-
// is already set to something other than the completion sentinel.
97-
98-
object oldContinuation = _continuation;
99-
if (oldContinuation == null)
100-
{
101-
_continuationState = state;
102-
oldContinuation = Interlocked.CompareExchange(ref _continuation, continuation, null);
103-
}
104-
105-
if (oldContinuation != null)
32+
try
10633
{
107-
// Operation already completed, so we need to call the supplied callback.
108-
if (!ReferenceEquals(oldContinuation, AutoResetValueTaskSourceCoreShared.s_sentinel))
109-
{
110-
throw new InvalidOperationException();
111-
}
112-
113-
continuation(state);
34+
return _sourceCore.GetResult(token);
11435
}
115-
}
116-
117-
private void ValidateToken(short token)
118-
{
119-
if (token != _version)
36+
finally
12037
{
121-
throw new InvalidOperationException();
38+
_sourceCore.Reset();
12239
}
12340
}
12441

125-
private void SignalCompletion()
126-
{
127-
if (_completed)
128-
{
129-
throw new InvalidOperationException();
130-
}
131-
_completed = true;
42+
void IValueTaskSource.GetResult(short token) => GetResult(token);
43+
TResult IValueTaskSource<TResult>.GetResult(short token) => GetResult(token);
13244

133-
Interlocked.CompareExchange(ref _continuation, AutoResetValueTaskSourceCoreShared.s_sentinel, null)
134-
?.Invoke(_continuationState);
135-
}
45+
ValueTaskSourceStatus IValueTaskSource.GetStatus(short token) => _sourceCore.GetStatus(token);
46+
ValueTaskSourceStatus IValueTaskSource<TResult>.GetStatus(short token) => _sourceCore.GetStatus(token);
13647

137-
void IValueTaskSource.GetResult(short token)
138-
{
139-
GetResult(token);
140-
}
141-
}
142-
143-
internal static class AutoResetValueTaskSourceCoreShared // separated out of generic to avoid unnecessary duplication
144-
{
145-
internal static readonly Action<object> s_sentinel = CompletionSentinel;
146-
private static void CompletionSentinel(object _) // named method to aid debugging
147-
{
148-
Debug.Fail("The sentinel delegate should never be invoked.");
149-
throw new InvalidOperationException();
150-
}
48+
void IValueTaskSource.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _sourceCore.OnCompleted(continuation, state, token, flags);
49+
void IValueTaskSource<TResult>.OnCompleted(Action<object> continuation, object state, short token, ValueTaskSourceOnCompletedFlags flags) => _sourceCore.OnCompleted(continuation, state, token, flags);
15150
}
15251
}

0 commit comments

Comments
 (0)