Skip to content

Commit 2458278

Browse files
committed
fix: WebSocket 연결 생명주기 및 메모리 누수 개선
문제점: - 1KB 작은 버퍼로 인한 성능 저하 - 장시간 연결에서 타임아웃 미처리 - 연결 해제 시 리소스 정리 불완전 - 하트비트 메커니즘 부재 수정사항: - 버퍼 크기 1KB → 4KB 증가로 성능 향상 - 30분 타임아웃으로 장시간 연결 관리 - ping/pong 하트비트 메시지 처리 추가 - 연결 종료 시 완전한 리소스 정리 - 초기 연결 확인 메시지 전송 - CancellationTokenSource 적절한 해제
1 parent c6da512 commit 2458278

File tree

1 file changed

+66
-8
lines changed

1 file changed

+66
-8
lines changed

ProjectVG.Api/Middleware/WebSocketMiddleware.cs

Lines changed: 66 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -103,28 +103,86 @@ private async Task RegisterConnection(Guid userId, WebSocket socket)
103103
await _webSocketService.ConnectAsync(userId.ToString());
104104
}
105105

106-
/// <summary>
107-
/// 세션 루프 실행
106+
/// <summary>
107+
/// 세션 루프 실행
108108
/// </summary>
109109
private async Task RunSessionLoop(WebSocket socket, string userId)
110110
{
111-
var buffer = new byte[1024];
111+
var buffer = new byte[1024 * 4]; // Increase buffer size for better performance
112+
var cancellationTokenSource = new CancellationTokenSource();
113+
114+
// Set a reasonable timeout for WebSocket operations
115+
cancellationTokenSource.CancelAfter(TimeSpan.FromMinutes(30));
116+
112117
try {
113-
while (socket.State == WebSocketState.Open) {
114-
var result = await socket.ReceiveAsync(new ArraySegment<byte>(buffer), CancellationToken.None);
118+
_logger.LogInformation("WebSocket 세션 시작: {UserId}", userId);
119+
120+
// Send initial connection confirmation
121+
var welcomeMessage = System.Text.Encoding.UTF8.GetBytes($"{{\"type\":\"connected\",\"userId\":\"{userId}\"}}");
122+
await socket.SendAsync(
123+
new ArraySegment<byte>(welcomeMessage),
124+
WebSocketMessageType.Text,
125+
true,
126+
cancellationTokenSource.Token).ConfigureAwait(false);
127+
128+
while (socket.State == WebSocketState.Open && !cancellationTokenSource.Token.IsCancellationRequested) {
129+
var result = await socket.ReceiveAsync(
130+
new ArraySegment<byte>(buffer),
131+
cancellationTokenSource.Token).ConfigureAwait(false);
115132

116133
if (result.MessageType == WebSocketMessageType.Close) {
117134
_logger.LogInformation("연결 종료 요청: {UserId}", userId);
118135
break;
119136
}
137+
138+
if (result.MessageType == WebSocketMessageType.Pong) {
139+
_logger.LogDebug("Pong 받음: {UserId}", userId);
140+
continue;
141+
}
142+
143+
// Handle heartbeat/ping messages
144+
if (result.MessageType == WebSocketMessageType.Text) {
145+
var message = System.Text.Encoding.UTF8.GetString(buffer, 0, result.Count);
146+
if (message.Contains("ping")) {
147+
var pongMessage = System.Text.Encoding.UTF8.GetBytes("{\"type\":\"pong\"}");
148+
await socket.SendAsync(
149+
new ArraySegment<byte>(pongMessage),
150+
WebSocketMessageType.Text,
151+
true,
152+
cancellationTokenSource.Token).ConfigureAwait(false);
153+
}
154+
}
120155
}
121156
}
157+
catch (OperationCanceledException) {
158+
_logger.LogWarning("WebSocket 세션 타임아웃: {UserId}", userId);
159+
}
160+
catch (WebSocketException ex) {
161+
_logger.LogWarning(ex, "WebSocket 연결 오류: {UserId}", userId);
162+
}
122163
catch (Exception ex) {
123-
_logger.LogError(ex, "세션 루프 오류: {UserId}", userId);
164+
_logger.LogError(ex, "세션 루프 예상치 못한 오류: {UserId}", userId);
124165
}
125166
finally {
126-
_logger.LogInformation("연결 해제: {UserId}", userId);
127-
await _webSocketService.DisconnectAsync(userId);
167+
_logger.LogInformation("WebSocket 연결 해제: {UserId}", userId);
168+
169+
try {
170+
await _webSocketService.DisconnectAsync(userId).ConfigureAwait(false);
171+
_connectionRegistry.Unregister(userId);
172+
173+
if (socket.State == WebSocketState.Open || socket.State == WebSocketState.CloseReceived) {
174+
await socket.CloseAsync(
175+
WebSocketCloseStatus.NormalClosure,
176+
"Connection closed",
177+
CancellationToken.None).ConfigureAwait(false);
178+
}
179+
}
180+
catch (Exception ex) {
181+
_logger.LogError(ex, "WebSocket 정리 중 오류: {UserId}", userId);
182+
}
183+
finally {
184+
cancellationTokenSource?.Dispose();
185+
}
128186
}
129187
}
130188
}

0 commit comments

Comments
 (0)