bool HandleWebsocketUpgrade(HttpSession session, int routeTableIndex, int endpointIndex, WebSocketUpgradeResponse response, out int responseBytesWritten) { responseBytesWritten = session.WriteWebSocketUpgradeResponse(response); var acceptUpgradeResponse = response as WebSocketUpgradeResponse.AcceptUpgradeResponse; if (acceptUpgradeResponse == null) { return false; } else { long id = Interlocked.Increment(ref acceptedWebSocketSessions); var webSocketSession = new WebSocketSession(id, session.TcpClient, session.Stream, bytesReceived => router.Metrics.CountRequestBytesIn(routeTableIndex, endpointIndex, bytesReceived), bytesSent => router.Metrics.CountResponseBytesOut(routeTableIndex, endpointIndex, bytesSent), () => UntrackWebSocketSession(id)); TrackWebSocketSession(webSocketSession); try { acceptUpgradeResponse.OnAccepted(webSocketSession); //TODO: Task.Run this? return true; } catch (Exception e) { UntrackWebSocketSession(id); logger.Error("Error calling WebSocketUpgradeResponse.OnAccepted callback - {0}", e.ToString()); return false; } } }
void TrackWebSocketSession(WebSocketSession session) { webSocketSessionTable[session.Id] = session; var sessionCount = webSocketSessionTable.Count; int currentMax; while ((currentMax = maxConnectedWebSocketSessions) < sessionCount) { if (Interlocked.CompareExchange(ref maxConnectedWebSocketSessions, sessionCount, currentMax) != currentMax) { continue; } } }
async Task<int> AcceptWebSocketUpgrade(HttpSession session, int routeTableIndex, int endpointIndex, WebSocketUpgradeResponse.AcceptUpgradeResponse response) { var bytesWritten = await response.WriteToAsync(session.Stream, 0, sessionReadTimeout).ConfigureAwait(false); var id = Interlocked.Increment(ref acceptedWebSocketSessions); var webSocketSession = new WebSocketSession(id, session.TcpClient, session.Stream, bytesReceived => router.Metrics.CountRequestBytesIn(routeTableIndex, endpointIndex, bytesReceived), bytesSent => router.Metrics.CountResponseBytesOut(routeTableIndex, endpointIndex, bytesSent), () => UntrackWebSocketSession(id), sessionReadBufferSize, (int)sessionReadTimeout.TotalMilliseconds, (int)sessionWriteTimeout.TotalMilliseconds); TrackWebSocketSession(webSocketSession); try { response.OnAccepted(webSocketSession); //TODO: Task.Run this? return bytesWritten; } catch (Exception) { UntrackWebSocketSession(id); throw; } }