Imported Upstream version 4.6.0.125

Former-commit-id: a2155e9bd80020e49e72e86c44da02a8ac0e57a4

mcs/class/referencesource/System.Web/WebSockets/AspNetWebSocketManager.cs

@@ -0,0 +1,87 @@
+//------------------------------------------------------------------------------
+// <copyright file="AspNetWebSocketManager.cs" company="Microsoft">
+//     Copyright (c) Microsoft Corporation.  All rights reserved.
+// </copyright>
+//------------------------------------------------------------------------------
+
+namespace System.Web.WebSockets {
+    using System;
+    using System.Collections.Generic;
+    using System.Linq;
+    using System.Threading.Tasks;
+    using System.Web.Util;
+
+    // Keeps track of AspNetWebSocket instances so that they can be aborted en masse,
+    // such as in the case of an AppDomain shutdown.
+
+    internal sealed class AspNetWebSocketManager {
+
+        public static readonly AspNetWebSocketManager Current = new AspNetWebSocketManager(PerfCounters.Instance);
+
+        private bool _aborted;
+        internal readonly HashSet<IAsyncAbortableWebSocket> _activeSockets = new HashSet<IAsyncAbortableWebSocket>(); // internal only for unit testing purposes
+        private readonly IPerfCounters _perfCounters;
+
+        internal AspNetWebSocketManager(IPerfCounters perfCounters) {
+            _perfCounters = perfCounters;
+        }
+
+        public int ActiveSocketCount {
+            get {
+                // We acquire a full lock when reading the count, similar to how the collections
+                // in the System.Collections.Concurrent namespace operate.
+                lock (_activeSockets) {
+                    return _activeSockets.Count;
+                }
+            }
+        }
+
+        // Calls Abort() on each tracked socket, then blocks until all have been aborted
+        public void AbortAllAndWait() {
+            // Make a copy so we're not iterating over the original collection asynchronously;
+            // keep the lock for as short a duration as possible.
+            IAsyncAbortableWebSocket[] sockets;
+            lock (_activeSockets) {
+                _aborted = true;
+                sockets = _activeSockets.ToArray();
+            }
+
+            Task[] abortTasks = Array.ConvertAll(sockets, socket => socket.AbortAsync());
+            Task.WaitAll(abortTasks);
+        }
+
+        // Begins tracking a socket, calling Abort() if there was an earlier call to AbortAll()
+        public void Add(IAsyncAbortableWebSocket webSocket) {
+            int activeSocketCount;
+            bool shouldAbort;
+
+            // keep the lock for as short a period as possible
+            lock (_activeSockets) {
+                _activeSockets.Add(webSocket);
+                activeSocketCount = _activeSockets.Count;
+                shouldAbort = _aborted;
+            }
+
+            // perform any additional operations outside the lock
+            _perfCounters.SetCounter(AppPerfCounter.REQUESTS_EXECUTING_WEBSOCKETS, activeSocketCount);
+            if (shouldAbort) {
+                webSocket.AbortAsync(); // don't care about the result of the abort at the present time
+            }
+        }
+
+        // Stops tracking a socket
+        public void Remove(IAsyncAbortableWebSocket webSocket) {
+            int activeSocketCount;
+
+            // keep the lock for as short a period as possible
+            lock (_activeSockets) {
+                _activeSockets.Remove(webSocket);
+                activeSocketCount = _activeSockets.Count;
+            }
+
+            // perform any additional operations outside the lock
+            _perfCounters.SetCounter(AppPerfCounter.REQUESTS_EXECUTING_WEBSOCKETS, activeSocketCount);
+        }
+
+    }
+}