|
2 | 2 | using System.Globalization; |
3 | 3 | using System.Runtime.ExceptionServices; |
4 | 4 | using System.Threading; |
| 5 | +using System.Threading.Tasks; |
5 | 6 |
|
6 | 7 | using Microsoft.Extensions.Logging; |
7 | 8 |
|
@@ -31,13 +32,8 @@ internal abstract class SubsystemSession : ISubsystemSession |
31 | 32 | private EventWaitHandle _channelClosedWaitHandle = new ManualResetEvent(initialState: false); |
32 | 33 | private bool _isDisposed; |
33 | 34 |
|
34 | | - /// <summary> |
35 | | - /// Gets or set the number of seconds to wait for an operation to complete. |
36 | | - /// </summary> |
37 | | - /// <value> |
38 | | - /// The number of seconds to wait for an operation to complete, or -1 to wait indefinitely. |
39 | | - /// </value> |
40 | | - public int OperationTimeout { get; private set; } |
| 35 | + /// <inheritdoc/> |
| 36 | + public int OperationTimeout { get; set; } |
41 | 37 |
|
42 | 38 | /// <summary> |
43 | 39 | /// Occurs when an error occurred. |
@@ -253,6 +249,59 @@ public void WaitOnHandle(WaitHandle waitHandle, int millisecondsTimeout) |
253 | 249 | } |
254 | 250 | } |
255 | 251 |
|
| 252 | + protected async Task<T> WaitOnHandleAsync<T>(TaskCompletionSource<T> tcs, int millisecondsTimeout, CancellationToken cancellationToken) |
| 253 | + { |
| 254 | + cancellationToken.ThrowIfCancellationRequested(); |
| 255 | + |
| 256 | + var errorOccuredReg = ThreadPool.RegisterWaitForSingleObject( |
| 257 | + _errorOccuredWaitHandle, |
| 258 | + (tcs, _) => ((TaskCompletionSource<T>)tcs).TrySetException(_exception), |
| 259 | + state: tcs, |
| 260 | + millisecondsTimeOutInterval: -1, |
| 261 | + executeOnlyOnce: true); |
| 262 | + |
| 263 | + var sessionDisconnectedReg = ThreadPool.RegisterWaitForSingleObject( |
| 264 | + _sessionDisconnectedWaitHandle, |
| 265 | + static (tcs, _) => ((TaskCompletionSource<T>)tcs).TrySetException(new SshException("Connection was closed by the server.")), |
| 266 | + state: tcs, |
| 267 | + millisecondsTimeOutInterval: -1, |
| 268 | + executeOnlyOnce: true); |
| 269 | + |
| 270 | + var channelClosedReg = ThreadPool.RegisterWaitForSingleObject( |
| 271 | + _channelClosedWaitHandle, |
| 272 | + static (tcs, _) => ((TaskCompletionSource<T>)tcs).TrySetException(new SshException("Channel was closed.")), |
| 273 | + state: tcs, |
| 274 | + millisecondsTimeOutInterval: -1, |
| 275 | + executeOnlyOnce: true); |
| 276 | + |
| 277 | + using var timeoutCts = new CancellationTokenSource(millisecondsTimeout); |
| 278 | + using var linkedCts = CancellationTokenSource.CreateLinkedTokenSource(cancellationToken, timeoutCts.Token); |
| 279 | + |
| 280 | + using var tokenReg = linkedCts.Token.Register( |
| 281 | + static s => |
| 282 | + { |
| 283 | + (var tcs, var cancellationToken) = ((TaskCompletionSource<T>, CancellationToken))s; |
| 284 | + _ = tcs.TrySetCanceled(cancellationToken); |
| 285 | + }, |
| 286 | + state: (tcs, cancellationToken), |
| 287 | + useSynchronizationContext: false); |
| 288 | + |
| 289 | + try |
| 290 | + { |
| 291 | + return await tcs.Task.ConfigureAwait(false); |
| 292 | + } |
| 293 | + catch (OperationCanceledException oce) when (timeoutCts.IsCancellationRequested) |
| 294 | + { |
| 295 | + throw new SshOperationTimeoutException("Operation has timed out.", oce); |
| 296 | + } |
| 297 | + finally |
| 298 | + { |
| 299 | + _ = errorOccuredReg.Unregister(waitObject: null); |
| 300 | + _ = sessionDisconnectedReg.Unregister(waitObject: null); |
| 301 | + _ = channelClosedReg.Unregister(waitObject: null); |
| 302 | + } |
| 303 | + } |
| 304 | + |
256 | 305 | /// <summary> |
257 | 306 | /// Blocks the current thread until the specified <see cref="WaitHandle"/> gets signaled, using a |
258 | 307 | /// 32-bit signed integer to specify the time interval in milliseconds. |
|
0 commit comments