forked from npgsql/npgsql
-
Notifications
You must be signed in to change notification settings - Fork 3
Expand file tree
/
Copy pathAsyncHelpers.cs
More file actions
143 lines (119 loc) · 5.87 KB
/
AsyncHelpers.cs
File metadata and controls
143 lines (119 loc) · 5.87 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
using System;
using System.Diagnostics;
using System.Runtime.CompilerServices;
using System.Threading;
using System.Threading.Tasks;
namespace Npgsql.Internal.Converters;
static class AsyncHelpers
{
public static void OnCompletedWithSource(Task task, CompletionSource source, CompletionSourceContinuation continuation)
{
_ = Core(task, source, continuation);
// Have our state machine be pooled, but don't return the task, source.Task should be used instead.
[AsyncMethodBuilder(typeof(PoolingAsyncValueTaskMethodBuilder))]
async ValueTask Core(Task task, CompletionSource source, CompletionSourceContinuation continuation)
{
try
{
await task.ConfigureAwait(false);
continuation.Invoke(task, source);
}
catch (Exception ex)
{
source.SetException(ex);
}
// Guarantee the type stays loaded until the function pointer call is done.
continuation.KeepAlive();
}
}
public abstract class CompletionSource
{
public abstract void SetException(Exception exception);
}
public sealed class CompletionSource<T> : CompletionSource
{
AsyncValueTaskMethodBuilder<T> _amb;
public ValueTask<T> Task { get; }
public CompletionSource()
{
_amb = AsyncValueTaskMethodBuilder<T>.Create();
// AsyncValueTaskMethodBuilder's Task and SetResult aren't thread safe in regard to each other
// Which is why we access it prematurely
Task = _amb.Task;
}
public void SetResult(T value)
=> _amb.SetResult(value);
public override void SetException(Exception exception)
=> _amb.SetException(exception);
}
public sealed class PoolingCompletionSource<T> : CompletionSource
{
PoolingAsyncValueTaskMethodBuilder<T> _amb;
public ValueTask<T> Task { get; }
public PoolingCompletionSource()
{
_amb = PoolingAsyncValueTaskMethodBuilder<T>.Create();
// PoolingAsyncValueTaskMethodBuilder's Task and SetResult aren't thread safe in regard to each other
// Which is why we access it prematurely
Task = _amb.Task;
}
public void SetResult(T value)
=> _amb.SetResult(value);
public override void SetException(Exception exception)
=> _amb.SetException(exception);
}
// Using a function pointer here is safe against assembly unloading as the instance reference that the static pointer method lives on is passed along.
// As such the instance cannot be collected by the gc which means the entire assembly is prevented from unloading until we're done.
// Split out into a struct as unsafe and async don't mix, while we do want a nicely typed function pointer signature to prevent mistakes.
public readonly unsafe struct CompletionSourceContinuation
{
readonly object _handle;
readonly delegate*<Task, CompletionSource, void> _continuation;
/// <param name="handle">A reference to the type that houses the static method <see ref="continuation"/> points to.</param>
/// <param name="continuation">The continuation</param>
public CompletionSourceContinuation(object handle, delegate*<Task, CompletionSource, void> continuation)
{
_handle = handle;
_continuation = continuation;
}
public void KeepAlive() => GC.KeepAlive(_handle);
public void Invoke(Task task, CompletionSource tcs) => _continuation(task, tcs);
}
public static unsafe ValueTask<T?> ReadAsyncAsNullable<T>(this PgConverter<T?> instance, PgConverter<T> effectiveConverter, PgReader reader, CancellationToken cancellationToken)
where T : struct
{
// Cheap if we have all the data.
var task = effectiveConverter.ReadAsync(reader, cancellationToken);
if (task.IsCompletedSuccessfully)
return new(new T?(task.Result));
// Otherwise we do one additional allocation, this allows us to share state machine codegen for all Ts.
var source = new PoolingCompletionSource<T?>();
OnCompletedWithSource(task.AsTask(), source, new(instance, &UnboxAndComplete));
return source.Task;
static void UnboxAndComplete(Task task, CompletionSource completionSource)
{
// Justification: exact type Unsafe.As used to reduce generic duplication cost.
Debug.Assert(task is Task<T>);
Debug.Assert(completionSource is PoolingCompletionSource<T?>);
Unsafe.As<PoolingCompletionSource<T?>>(completionSource).SetResult(new T?(new ValueTask<T>(Unsafe.As<Task<T>>(task)).Result));
}
}
public static unsafe ValueTask<T> ReadAsObjectAsyncAsT<T>(this PgConverter<T> instance, PgConverter effectiveConverter, PgReader reader, CancellationToken cancellationToken)
{
// Cheap if we have all the data.
var task = effectiveConverter.ReadAsObjectAsync(reader, cancellationToken);
if (task.IsCompletedSuccessfully)
return new((T)task.Result);
// Otherwise we do one additional allocation, this allows us to share state machine codegen for all Ts.
var source = new PoolingCompletionSource<T>();
OnCompletedWithSource(task.AsTask(), source, new(instance, &UnboxAndComplete));
return source.Task;
static void UnboxAndComplete(Task task, CompletionSource completionSource)
{
// Justification: exact type Unsafe.As used to reduce generic duplication cost.
Debug.Assert(task is Task<object>);
Debug.Assert(completionSource is PoolingCompletionSource<T>);
Unsafe.As<PoolingCompletionSource<T>>(completionSource).SetResult((T)new ValueTask<object>(Unsafe.As<Task<object>>(task)).Result);
}
}
}