-
Notifications
You must be signed in to change notification settings - Fork 9.9k
/
TransportConnectionManager.cs
103 lines (85 loc) · 3.56 KB
/
TransportConnectionManager.cs
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
// Copyright (c) .NET Foundation. All rights reserved.
// Licensed under the Apache License, Version 2.0. See License.txt in the project root for license information.
#nullable enable
using System;
using System.Collections.Concurrent;
using System.Collections.Generic;
using System.Threading;
using System.Threading.Tasks;
using Microsoft.AspNetCore.Connections;
namespace Microsoft.AspNetCore.Server.Kestrel.Core.Internal.Infrastructure
{
internal class TransportConnectionManager
{
private readonly ConnectionManager _connectionManager;
private readonly ConcurrentDictionary<long, ConnectionReference> _connectionReferences = new ConcurrentDictionary<long, ConnectionReference>();
public TransportConnectionManager(ConnectionManager connectionManager)
{
_connectionManager = connectionManager;
}
public void AddConnection(long id, KestrelConnection connection)
{
var connectionReference = new ConnectionReference(id, connection, this);
if (!_connectionReferences.TryAdd(id, connectionReference))
{
throw new ArgumentException(nameof(id));
}
_connectionManager.AddConnection(id, connectionReference);
}
public void RemoveConnection(long id)
{
if (!_connectionReferences.TryRemove(id, out _))
{
throw new ArgumentException(nameof(id));
}
_connectionManager.RemoveConnection(id);
}
// This is only called by the ConnectionManager when the connection reference becomes
// unrooted because the application never completed.
public void StopTracking(long id)
{
if (!_connectionReferences.TryRemove(id, out _))
{
throw new ArgumentException(nameof(id));
}
}
public async Task<bool> CloseAllConnectionsAsync(CancellationToken token)
{
var closeTasks = new List<Task>();
foreach (var kvp in _connectionReferences)
{
if (kvp.Value.TryGetConnection(out var connection))
{
connection.RequestClose();
closeTasks.Add(connection.ExecutionTask);
}
}
var allClosedTask = Task.WhenAll(closeTasks.ToArray());
return await Task.WhenAny(allClosedTask, CancellationTokenAsTask(token)).ConfigureAwait(false) == allClosedTask;
}
public async Task<bool> AbortAllConnectionsAsync()
{
var abortTasks = new List<Task>();
foreach (var kvp in _connectionReferences)
{
if (kvp.Value.TryGetConnection(out var connection))
{
connection.TransportConnection.Abort(new ConnectionAbortedException(CoreStrings.ConnectionAbortedDuringServerShutdown));
abortTasks.Add(connection.ExecutionTask);
}
}
var allAbortedTask = Task.WhenAll(abortTasks.ToArray());
return await Task.WhenAny(allAbortedTask, Task.Delay(1000)).ConfigureAwait(false) == allAbortedTask;
}
private static Task CancellationTokenAsTask(CancellationToken token)
{
if (token.IsCancellationRequested)
{
return Task.CompletedTask;
}
var tcs = new TaskCompletionSource(TaskCreationOptions.RunContinuationsAsynchronously);
token.Register(() => tcs.SetResult());
return tcs.Task;
}
}
}