Repository navigation
Expand file tree
/
Copy pathVirtualConnectionClient.cs
More file actions
238 lines (206 loc) · 9.38 KB
/
Copy pathVirtualConnectionClient.cs
File metadata and controls
238 lines (206 loc) · 9.38 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
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
using System;
using System.Collections.Generic;
using System.IO;
using System.Linq;
using System.Threading;
using System.Threading.Tasks;
namespace Microsoft.AspNetCore.NodeServices.Sockets.VirtualConnections
{
/// <summary>
/// A callback that will be invoked if the <see cref="VirtualConnectionClient"/> encounters a read error.
/// </summary>
/// <param name="ex"></param>
public delegate void VirtualConnectionReadErrorHandler(Exception ex);
/// <summary>
/// Wraps an underlying physical read/write stream (e.g., named pipes, domain sockets, or TCP sockets) and
/// exposes an API for making 'virtual connections', which act as independent read/write streams.
/// Traffic over these virtual connections is multiplexed over the underlying physical stream. This is useful
/// for fast stream-based inter-process communication because it avoids the overhead of opening a new physical
/// connection each time a new communication channel is needed.
/// </summary>
internal class VirtualConnectionClient : IDisposable
{
internal const int MaxFrameBodyLength = 16 * 1024;
public event VirtualConnectionReadErrorHandler OnError;
private Stream _underlyingTransport;
private Dictionary<long, VirtualConnection> _activeInnerStreams;
private long _nextInnerStreamId;
private readonly SemaphoreSlim _streamWriterSemaphore = new SemaphoreSlim(1);
private readonly object _readControlLock = new object();
private Exception _readLoopExitedWithException;
private readonly CancellationTokenSource _disposalCancellatonToken = new CancellationTokenSource();
private bool _disposedValue = false;
public VirtualConnectionClient(Stream underlyingTransport)
{
_underlyingTransport = underlyingTransport;
_activeInnerStreams = new Dictionary<long, VirtualConnection>();
RunReadLoop();
}
public Stream OpenVirtualConnection()
{
// Improve discoverability of read-loop errors (in case the developer doesn't add an OnError listener)
ThrowIfReadLoopFailed();
var id = Interlocked.Increment(ref _nextInnerStreamId);
var newInnerStream = new VirtualConnection(id, this);
lock (_activeInnerStreams)
{
_activeInnerStreams.Add(id, newInnerStream);
}
return newInnerStream;
}
// It's async void because nothing waits for it to finish (it continues indefinitely). It signals any errors via
// a separate channel.
private async void RunReadLoop()
{
try
{
while (!_disposalCancellatonToken.IsCancellationRequested)
{
var remoteIsStillConnected = await ProcessNextFrameAsync();
if (!remoteIsStillConnected)
{
CloseAllActiveStreams();
}
}
}
catch (Exception ex)
{
// Not all underlying transports correctly honor cancellation tokens. For example,
// DomainSocketStreamTransport's ReadAsync ignores them, so we only know to stop
// the read loop when the underlying stream is disposed and then it throws ObjectDisposedException.
if (!(ex is TaskCanceledException || ex is ObjectDisposedException))
{
_readLoopExitedWithException = ex;
var evt = OnError;
if (evt != null)
{
evt(ex);
}
}
}
}
private async Task<bool> ProcessNextFrameAsync()
{
// First read frame header
var frameHeaderBuffer = await ReadExactLength(12);
if (frameHeaderBuffer == null)
{
return false; // Underlying stream was closed
}
// Parse frame header, then read the frame body
long streamId = BitConverter.ToInt64(frameHeaderBuffer, 0);
int frameBodyLength = BitConverter.ToInt32(frameHeaderBuffer, 8);
if (frameBodyLength < 0 || frameBodyLength > MaxFrameBodyLength)
{
throw new InvalidDataException("Illegal frame length: " + frameBodyLength);
}
var frameBody = await ReadExactLength(frameBodyLength);
if (frameBody == null)
{
return false; // Underlying stream was closed
}
// Dispatch the frame to the relevant inner stream
VirtualConnection innerStream;
lock (_activeInnerStreams)
{
_activeInnerStreams.TryGetValue(streamId, out innerStream);
}
if (innerStream != null)
{
await innerStream.AddDataToQueue(frameBody);
}
return true;
}
private async Task<byte[]> ReadExactLength(int lengthToRead) {
byte[] buffer = new byte[lengthToRead];
var totalBytesRead = 0;
var ct = _disposalCancellatonToken.Token;
while (totalBytesRead < lengthToRead)
{
var chunkLengthRead = await _underlyingTransport.ReadAsync(buffer, totalBytesRead, lengthToRead - totalBytesRead, ct);
if (chunkLengthRead == 0)
{
// Underlying stream was closed
return null;
}
totalBytesRead += chunkLengthRead;
}
return buffer;
}
private void CloseAllActiveStreams()
{
IList<VirtualConnection> innerStreamsCopy;
// Only hold the lock while cloning the list of inner streams. Release the lock before
// actually disposing them, because each 'dispose' call will try to take another lock
// so it can remove that inner stream from activeInnerStreams.
lock (_activeInnerStreams)
{
innerStreamsCopy = _activeInnerStreams.Values.ToList();
}
foreach (var stream in innerStreamsCopy)
{
stream.Dispose();
}
}
public void Dispose()
{
if (!_disposedValue)
{
_disposedValue = true;
_disposalCancellatonToken.Cancel(); // Stops the read loop
CloseAllActiveStreams();
}
}
public async Task WriteAsync(long innerStreamId, byte[] data, int offset, int count, CancellationToken cancellationToken)
{
// In case the amount of data to be sent exceeds the max frame length, split it into separate frames
// Note that we always send at least one frame, even if it's empty, because the zero-length frame is the signal to close a virtual connection
// (hence 'do..while' instead of just 'while').
int bytesWritten = 0;
do {
// Improve discoverability of read-loop errors (in case the developer doesn't add an OnError listener)
ThrowIfReadLoopFailed();
// Hold the write lock only for the time taken to send a single frame, not all frames, to allow large sends to be proceed in parallel
await _streamWriterSemaphore.WaitAsync(cancellationToken).ConfigureAwait(false);
try
{
// Write stream ID, then length prefix, then chunk payload, then flush
var nextChunkBodyLength = Math.Min(MaxFrameBodyLength, count - bytesWritten);
await _underlyingTransport.WriteAsync(BitConverter.GetBytes(innerStreamId), 0, 8, cancellationToken).ConfigureAwait(false);
await _underlyingTransport.WriteAsync(BitConverter.GetBytes(nextChunkBodyLength), 0, 4, cancellationToken).ConfigureAwait(false);
if (nextChunkBodyLength > 0)
{
await _underlyingTransport.WriteAsync(data, offset + bytesWritten, nextChunkBodyLength, cancellationToken).ConfigureAwait(false);
bytesWritten += nextChunkBodyLength;
}
await _underlyingTransport.FlushAsync(cancellationToken).ConfigureAwait(false);
}
finally
{
_streamWriterSemaphore.Release();
}
} while (bytesWritten < count);
}
public void CloseInnerStream(long innerStreamId, bool isAlreadyClosedRemotely)
{
lock (_activeInnerStreams)
{
if (_activeInnerStreams.ContainsKey(innerStreamId))
{
_activeInnerStreams.Remove(innerStreamId);
}
}
if (!isAlreadyClosedRemotely) {
// Also notify the remote that this innerstream is closed
WriteAsync(innerStreamId, new byte[0], 0, 0, new CancellationToken()).Wait();
}
}
private void ThrowIfReadLoopFailed()
{
if (_readLoopExitedWithException != null)
{
throw new AggregateException("The connection failed - see InnerException for details.", _readLoopExitedWithException);
}
}
}
}