Skip to content

Commit 129adf6

Browse files
committed
Harden WebSocket reconnection logic in Unity plugin
Three defensive fixes to prevent stale state from blocking reconnection: - Add _connecting guard to prevent overlapping ConnectAsync calls - Check socket identity before clobbering _isConnected in catch/finally blocks, so a stale receive loop or failed attempt can't reset state that a newer successful connection already set - Add 10s timeout on ConnectAsync so hanging TCP connections don't block reconnection indefinitely
1 parent 402eda3 commit 129adf6

1 file changed

Lines changed: 31 additions & 6 deletions

File tree

UnityCtl.UnityPackage/Editor/UnityCtlClient.cs

Lines changed: 31 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -34,6 +34,7 @@ public class UnityCtlClient
3434
private float _lastConnectionCheck;
3535
private float _lastReconnectAttempt;
3636
private int _reconnectAttempts;
37+
private bool _connecting;
3738

3839
// Public connection status
3940
public bool IsConnected
@@ -139,6 +140,15 @@ private async Task ConnectAsync(int port)
139140
return;
140141
}
141142

143+
// Skip if another connection attempt is already in-flight
144+
if (_connecting)
145+
{
146+
DebugLog("[UnityCtl] Connection attempt already in progress, skipping");
147+
return;
148+
}
149+
150+
_connecting = true;
151+
142152
oldSocket = _webSocket;
143153
newSocket = new ClientWebSocket();
144154
_webSocket = newSocket;
@@ -165,11 +175,16 @@ private async Task ConnectAsync(int port)
165175

166176
try
167177
{
168-
await newSocket.ConnectAsync(uri, newCts.Token);
178+
// Timeout the connection attempt so it can't hang indefinitely
179+
using var connectCts = CancellationTokenSource.CreateLinkedTokenSource(newCts.Token);
180+
connectCts.CancelAfter(TimeSpan.FromSeconds(10));
181+
182+
await newSocket.ConnectAsync(uri, connectCts.Token);
169183

170184
lock (_connectionLock)
171185
{
172186
_isConnected = true;
187+
_connecting = false;
173188
}
174189
_reconnectAttempts = 0; // Reset reconnection counter on successful connection
175190

@@ -181,15 +196,21 @@ private async Task ConnectAsync(int port)
181196
// Flush any logs that were buffered before connection
182197
FlushLogBuffer();
183198

184-
// Start receive loop
185-
_ = Task.Run(() => ReceiveLoopAsync(newCts.Token));
199+
// Start receive loop (pass socket ref so finally block can check staleness)
200+
var socketRef = newSocket;
201+
_ = Task.Run(() => ReceiveLoopAsync(socketRef, newCts.Token));
186202
}
187203
catch (Exception ex)
188204
{
189205
DebugLogWarning($"[UnityCtl] ✗ Connection failed: {ex.Message}");
190206
lock (_connectionLock)
191207
{
192-
_isConnected = false;
208+
// Only update state if we're still the current attempt
209+
if (_webSocket == newSocket)
210+
{
211+
_isConnected = false;
212+
}
213+
_connecting = false;
193214
}
194215
}
195216
}
@@ -228,7 +249,7 @@ private static string GetPluginVersion()
228249
return "unknown";
229250
}
230251

231-
private async Task ReceiveLoopAsync(CancellationToken cancellationToken)
252+
private async Task ReceiveLoopAsync(ClientWebSocket ownSocket, CancellationToken cancellationToken)
232253
{
233254
var buffer = new byte[1024 * 16];
234255
var messageBuilder = new System.Text.StringBuilder();
@@ -286,7 +307,11 @@ private async Task ReceiveLoopAsync(CancellationToken cancellationToken)
286307
{
287308
lock (_connectionLock)
288309
{
289-
_isConnected = false;
310+
// Only clear connected state if we're still the active connection
311+
if (_webSocket == ownSocket)
312+
{
313+
_isConnected = false;
314+
}
290315
}
291316
DebugLog("[UnityCtl] Disconnected from bridge - will attempt to reconnect");
292317
}

0 commit comments

Comments
 (0)