Skip to content

Commit 9a8b603

Browse files
rekhoffjdetter
andauthored
Adding Abort to C# Websocket (#3352)
# Description of Changes The implementation of a solution to #3044 , this adds an `Abort` function to the `WebSocket`, which runs if `Disconnect` is called when the `WebSocket` is not connected. # API and ABI breaking changes Not API breaking. # Expected complexity level and risk 1 # Testing - [X] Test locally with a C# CLI test client. **Note**: Before change (either on Rust of C# server), server would see 4 `Debug` log entries about connecting, but not see the `Info` log about the client connection ending like would normally be seen in a disconnect. After change, server shows no log entries at all, because connection is properly aborted. - [x] Test locally with a C# WebGL test client. --------- Signed-off-by: rekhoff <r.ekhoff@clockworklabs.io> Co-authored-by: John Detter <4099508+jdetter@users.noreply.github.com>
1 parent 6772f0c commit 9a8b603

2 files changed

Lines changed: 88 additions & 4 deletions

File tree

sdks/csharp/src/SpacetimeDBClient.cs

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -508,6 +508,14 @@ public void Disconnect()
508508
{
509509
webSocket.Close();
510510
}
511+
#if UNITY_WEBGL && !UNITY_EDITOR
512+
else if (webSocket.IsConnecting)
513+
#else
514+
else if (webSocket.IsConnecting || webSocket.IsNoneState)
515+
#endif
516+
{
517+
webSocket.Abort(); // forceful during connecting
518+
}
511519

512520
_parseCancellationTokenSource.Cancel();
513521
}

sdks/csharp/src/WebSocket.cs

Lines changed: 80 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -38,6 +38,7 @@ public struct ConnectOptions
3838
private readonly ConcurrentQueue<Action> dispatchQueue = new();
3939

4040
protected ClientWebSocket Ws = new();
41+
private CancellationTokenSource? _connectCts;
4142

4243
public WebSocket(ConnectOptions options)
4344
{
@@ -60,9 +61,13 @@ public WebSocket(ConnectOptions options)
6061
#if UNITY_WEBGL && !UNITY_EDITOR
6162
private bool _isConnected = false;
6263
private bool _isConnecting = false;
64+
private bool _cancelConnectRequested = false;
6365
public bool IsConnected => _isConnected;
64-
#else
66+
public bool IsConnecting => _isConnecting;
67+
#else
6568
public bool IsConnected { get { return Ws != null && Ws.State == WebSocketState.Open; } }
69+
public bool IsConnecting { get { return Ws != null && Ws.State == WebSocketState.Connecting; } }
70+
public bool IsNoneState { get { return Ws != null && Ws.State == WebSocketState.None; } }
6671
#endif
6772

6873
#if UNITY_WEBGL && !UNITY_EDITOR
@@ -147,6 +152,7 @@ public async Task Connect(string? auth, string host, string nameOrAddress, Conne
147152
if (_isConnecting || _isConnected) return;
148153

149154
_isConnecting = true;
155+
_cancelConnectRequested = false;
150156
try
151157
{
152158
var uri = $"{host}/v1/database/{nameOrAddress}/subscribe?connection_id={connectionId}&compression={compression}";
@@ -167,6 +173,11 @@ public async Task Connect(string? auth, string host, string nameOrAddress, Conne
167173
dispatchQueue.Enqueue(() => OnConnectError?.Invoke(
168174
new Exception("Failed to connect WebSocket")));
169175
}
176+
else if (_cancelConnectRequested)
177+
{
178+
// If cancel was requested before open, proactively close now.
179+
WebSocket_Close(_webglSocketId, (int)WebSocketCloseStatus.NormalClosure, "Canceled during connect.");
180+
}
170181
}
171182
catch (Exception e)
172183
{
@@ -192,7 +203,7 @@ public async Task Connect(string? auth, string host, string nameOrAddress, Conne
192203
var url = new Uri(uri);
193204
Ws.Options.AddSubProtocol(_options.Protocol);
194205

195-
var source = new CancellationTokenSource(10000);
206+
_connectCts = new CancellationTokenSource(10000);
196207
if (!string.IsNullOrEmpty(auth))
197208
{
198209
Ws.Options.SetRequestHeader("Authorization", $"Bearer {auth}");
@@ -204,7 +215,7 @@ public async Task Connect(string? auth, string host, string nameOrAddress, Conne
204215

205216
try
206217
{
207-
await Ws.ConnectAsync(url, source.Token);
218+
await Ws.ConnectAsync(url, _connectCts.Token);
208219
if (Ws.State == WebSocketState.Open)
209220
{
210221
if (OnConnect != null)
@@ -376,14 +387,35 @@ await Ws.CloseAsync(WebSocketCloseStatus.MessageTooBig, closeMessage,
376387
#endif
377388
}
378389

390+
/// <summary>
391+
/// Cancel an in-flight ConnectAsync. Safe to call if no connect is pending.
392+
/// </summary>
393+
public void CancelConnect()
394+
{
395+
#if UNITY_WEBGL && !UNITY_EDITOR
396+
// No CTS on WebGL. Mark cancel intent so that when socket id arrives or open fires,
397+
// we immediately close and avoid reporting a connected state.
398+
_cancelConnectRequested = true;
399+
#else
400+
try { _connectCts?.Cancel(); } catch { /* ignore */ }
401+
#endif
402+
}
403+
379404
public Task Close(WebSocketCloseStatus code = WebSocketCloseStatus.NormalClosure)
380405
{
381406
#if UNITY_WEBGL && !UNITY_EDITOR
382-
if (_isConnected && _webglSocketId >= 0)
407+
if (_webglSocketId >= 0)
383408
{
409+
// If connected or connecting with a valid socket id, request a close.
384410
WebSocket_Close(_webglSocketId, (int)code, "Disconnecting normally.");
411+
_cancelConnectRequested = false; // graceful close intent
385412
_isConnected = false;
386413
}
414+
else if (_isConnecting)
415+
{
416+
// We don't yet have a socket id; remember to cancel once it arrives/opens.
417+
_cancelConnectRequested = true;
418+
}
387419
#else
388420
if (Ws?.State == WebSocketState.Open)
389421
{
@@ -393,6 +425,35 @@ public Task Close(WebSocketCloseStatus code = WebSocketCloseStatus.NormalClosure
393425
return Task.CompletedTask;
394426
}
395427

428+
/// <summary>
429+
/// Forcefully abort the WebSocket connection. This terminates any in-flight connect/receive/send
430+
/// and ensures the server-side socket is torn down promptly. Prefer Close() for graceful shutdowns.
431+
/// </summary>
432+
public void Abort()
433+
{
434+
#if UNITY_WEBGL && !UNITY_EDITOR
435+
if (_webglSocketId >= 0)
436+
{
437+
WebSocket_Close(_webglSocketId, (int)WebSocketCloseStatus.NormalClosure, "Aborting connection.");
438+
_isConnected = false;
439+
}
440+
else if (_isConnecting)
441+
{
442+
// No socket yet; ensure we close immediately once it opens.
443+
_cancelConnectRequested = true;
444+
}
445+
#else
446+
try
447+
{
448+
Ws?.Abort();
449+
}
450+
catch
451+
{
452+
// Intentionally swallow; Abort is best-effort.
453+
}
454+
#endif
455+
}
456+
396457
private Task? senderTask;
397458
private readonly ConcurrentQueue<ClientMessage> messageSendQueue = new();
398459

@@ -459,11 +520,21 @@ public WebSocketState GetState()
459520
{
460521
return Ws!.State;
461522
}
523+
462524
#if UNITY_WEBGL && !UNITY_EDITOR
463525
public void HandleWebGLOpen(int socketId)
464526
{
465527
if (socketId == _webglSocketId)
466528
{
529+
if (_cancelConnectRequested)
530+
{
531+
// Immediately close instead of reporting connected.
532+
WebSocket_Close(_webglSocketId, (int)WebSocketCloseStatus.NormalClosure, "Canceled during connect.");
533+
_isConnecting = false;
534+
_isConnected = false;
535+
_cancelConnectRequested = false;
536+
return;
537+
}
467538
_isConnected = true;
468539
if (OnConnect != null)
469540
dispatchQueue.Enqueue(() => OnConnect());
@@ -484,6 +555,9 @@ public void HandleWebGLClose(int socketId, int code, string reason)
484555
if (socketId == _webglSocketId && OnClose != null)
485556
{
486557
_isConnected = false;
558+
_isConnecting = false;
559+
_webglSocketId = -1;
560+
_cancelConnectRequested = false;
487561
var ex = code != (int)WebSocketCloseStatus.NormalClosure ? new Exception($"WebSocket closed with code {code}: {reason}") : null;
488562
dispatchQueue.Enqueue(() => OnClose?.Invoke(ex));
489563
}
@@ -494,6 +568,8 @@ public void HandleWebGLError(int socketId)
494568
UnityEngine.Debug.Log($"HandleWebGLError: {socketId}");
495569
if (socketId == _webglSocketId && OnConnectError != null)
496570
{
571+
_isConnecting = false;
572+
_webglSocketId = -1;
497573
dispatchQueue.Enqueue(() => OnConnectError(new Exception($"Socket {socketId} error.")));
498574
}
499575
}

0 commit comments

Comments
 (0)