fix: maybe deadlock in coordinator

BUAA-SE-Compiling · Oct 1, 2021 · ba2d0bc · ba2d0bc
1 parent 7dd9b2b
commit ba2d0bc
Showing 1 changed file with 20 additions and 21 deletions.
diff --git 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;
+            }
         }
     }