From ba2d0bcc35c5017c70851483c6b40d0e14e10f53 Mon Sep 17 00:00:00 2001
From: Rynco Maekawa <lynzrand@outlook.com>
Date: Fri, 1 Oct 2021 11:48:04 +0800
Subject: [PATCH] fix: maybe deadlock in coordinator

---
 .../Services/JudgerCoordinatorService.cs      | 41 +++++++++----------
 1 file changed, 20 insertions(+), 21 deletions(-)

diff --git a/coordinator/Services/JudgerCoordinatorService.cs b/coordinator/Services/JudgerCoordinatorService.cs
index 075c6a8d..09ac9c3c 100644
--- a/coordinator/Services/JudgerCoordinatorService.cs
+++ b/coordinator/Services/JudgerCoordinatorService.cs
@@ -105,19 +105,18 @@ public async ValueTask<bool> TryUseConnection(HttpContext ctx) {
                     ctx.Request.Query.TryGetValue("conn", out var connId_);
                     var connId = connId_.FirstOrDefault();
 
-                    var connLock = await connectionLock.LockAsync();
-                    if (connections.TryGetValue(auth, out var lastConn)) {
-                        if (lastConn.ConnectionId != null && connId != null && lastConn.ConnectionId == connId) {
-                            // replace this session
-                            await lastConn.Socket.Close(System.Net.WebSockets.WebSocketCloseStatus.PolicyViolation, "Duplicate connection", CancellationToken.None);
-                            connections.Remove(auth);
-                        } else {
-                            ctx.Response.StatusCode = StatusCodes.Status409Conflict;
-                            connLock.Dispose();
-                            return false;
+                    using (await connectionLock.LockAsync()) {
+                        if (connections.TryGetValue(auth, out var lastConn)) {
+                            if (lastConn.ConnectionId != null && connId != null && lastConn.ConnectionId == connId) {
+                                // replace this session
+                                await lastConn.Socket.Close(System.Net.WebSockets.WebSocketCloseStatus.PolicyViolation, "Duplicate connection", CancellationToken.None);
+                                connections.Remove(auth);
+                            } else {
+                                ctx.Response.StatusCode = StatusCodes.Status409Conflict;
+                                return false;
+                            }
                         }
                     }
-                    connLock.Dispose();
 
                     var ws = await ctx.WebSockets.AcceptWebSocketAsync();
 
@@ -601,16 +600,16 @@ public void Stop() {
         /// </summary>
         /// <returns>(connectedCount, runningCount)</returns>
         public async Task<(int, int)> GetConnectedJudgerInfo() {
-            await connectionLock.WaitAsync();
-            var result = connections.Aggregate((0, 0), (cnt, val) => {
-                if (val.Value.ActiveTaskCount > 0) {
-                    return (cnt.Item1 + 1, cnt.Item2 + 1);
-                } else {
-                    return (cnt.Item1 + 1, cnt.Item2);
-                }
-            });
-            connectionLock.Release();
-            return result;
+            using (await connectionLock.LockAsync()) {
+                var result = connections.Aggregate((0, 0), (cnt, val) => {
+                    if (val.Value.ActiveTaskCount > 0) {
+                        return (cnt.Item1 + 1, cnt.Item2 + 1);
+                    } else {
+                        return (cnt.Item1 + 1, cnt.Item2);
+                    }
+                });
+                return result;
+            }
         }
     }