-
Notifications
You must be signed in to change notification settings - Fork 48
/
LavaSocket.cs
197 lines (167 loc) · 7.14 KB
/
LavaSocket.cs
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
using System;
using System.Collections.Generic;
using System.Net.WebSockets;
using System.Text.Json;
using System.Threading;
using System.Threading.Tasks;
namespace Victoria {
/// <summary>
/// Wrapper around <see cref="ClientWebSocket"/> to make ws connections easier to handle.
/// </summary>
public sealed class LavaSocket : IAsyncDisposable {
private readonly IDictionary<string, string> _headers;
private readonly LavaConfig _lavaConfig;
private int _connectionAttempts;
private TimeSpan _reconnectInterval;
private bool _refIsUsable;
private ClientWebSocket _socket;
private CancellationTokenSource _tokenSource;
/// <summary>
/// Fires when connection is lost and a retry attempt is made.
/// </summary>
public event Func<string, Task> OnRetry;
/// <summary>
/// Fires when connection is established.
/// </summary>
public event Func<Task> OnConnected;
/// <summary>
/// Fires when either client or server closes connection.
/// </summary>
public event Func<string, Task> OnDisconnected;
/// <summary>
/// Fires when data is received from server.
/// </summary>
public event Func<byte[], Task> OnReceive;
internal LavaSocket(LavaConfig lavaConfig) {
_lavaConfig = lavaConfig;
_headers = new Dictionary<string, string>(3);
}
/// <inheritdoc />
public async ValueTask DisposeAsync() {
if (_socket.State == WebSocketState.Open) {
await _socket.CloseAsync(WebSocketCloseStatus.Empty, "", _tokenSource.Token)
.ConfigureAwait(false);
}
_tokenSource.CancelAfter(TimeSpan.FromSeconds(5));
_socket.Dispose();
}
/// <summary>
/// Set websocket headers before making a handshake request.
/// </summary>
/// <param name="key"></param>
/// <param name="value"></param>
public void SetHeader(string key, string value) {
if (_headers.ContainsKey(key)) {
return;
}
_headers.Add(key, value);
}
/// <summary>
/// Initializes instance of <see cref="ClientWebSocket"/> and connects to server.
/// </summary>
public async Task ConnectAsync() {
_tokenSource = new CancellationTokenSource();
_socket = new ClientWebSocket();
foreach (var (key, value) in _headers) {
_socket.Options.SetRequestHeader(key, value);
}
var url = new Uri($"ws://{_lavaConfig.Hostname}:{_lavaConfig.Port}");
if (_connectionAttempts == _lavaConfig.ReconnectAttempts) {
return;
}
try {
await _socket.ConnectAsync(url, CancellationToken.None).ContinueWith(VerifyConnectionAsync);
}
catch {
// IGNORE
}
}
/// <summary>
/// Used to send data to websocket server.
/// </summary>
/// <param name="value">Value of <typeparamref name="T"/></param>
/// <typeparam name="T">Type of data to set.</typeparam>
/// <exception cref="InvalidOperationException">Throws if not connected to websocket.</exception>
public async Task SendAsync<T>(T value) {
if (_socket.State != WebSocketState.Open) {
throw new InvalidOperationException("WebSocket connection to server isn't in open state.");
}
var rawBytes = JsonSerializer.SerializeToUtf8Bytes(value);
await _socket.SendAsync(rawBytes, WebSocketMessageType.Text, true, _tokenSource.Token)
.ConfigureAwait(false);
}
/// <summary>
/// Closes connection to websocket server.
/// </summary>
/// <returns></returns>
public async Task DisconnectAsync() {
await _socket.CloseAsync(WebSocketCloseStatus.NormalClosure, "Requested.", _tokenSource.Token)
.ConfigureAwait(false);
}
private async Task VerifyConnectionAsync(Task task) {
if (task.IsCanceled || task.IsFaulted || task.Exception != null) {
Volatile.Write(ref _refIsUsable, false);
await RetryConnectionAsync();
}
else {
Volatile.Write(ref _refIsUsable, true);
_connectionAttempts = 0;
OnConnected?.Invoke();
await ReceiveAsync()
.ConfigureAwait(false);
}
}
private async Task RetryConnectionAsync() {
_tokenSource.Cancel(false);
if (_connectionAttempts > _lavaConfig.ReconnectAttempts && _lavaConfig.ReconnectAttempts != -1) {
return;
}
if (Volatile.Read(ref _refIsUsable)) {
return;
}
Interlocked.Increment(ref _connectionAttempts);
_reconnectInterval += _lavaConfig.ReconnectDelay;
OnRetry?.Invoke(_connectionAttempts == _lavaConfig.ReconnectAttempts
? "This was the last attempt at re-establishing websocket connection."
: $"Waiting {_reconnectInterval.TotalSeconds}s before attempt #{_connectionAttempts}.");
await Task.Delay(_reconnectInterval)
.ContinueWith(_ => ConnectAsync())
.ConfigureAwait(false);
}
private async Task ReceiveAsync() {
try {
while (Volatile.Read(ref _refIsUsable) && _socket.State == WebSocketState.Open &&
!_tokenSource.IsCancellationRequested) {
var buffer = new byte[_lavaConfig.BufferSize];
var result = await _socket.ReceiveAsync(buffer, _tokenSource.Token)
.ConfigureAwait(false);
switch (result.MessageType) {
case WebSocketMessageType.Close:
Volatile.Write(ref _refIsUsable, false);
OnDisconnected?.Invoke("Server closed the connection!");
await RetryConnectionAsync()
.ConfigureAwait(false);
break;
case WebSocketMessageType.Text:
if (!result.EndOfMessage) {
continue;
}
var startLength = buffer.Length - 1;
while (buffer[startLength] == 0) {
--startLength;
}
var cleaned = new byte[startLength + 1];
Array.Copy(buffer, cleaned, startLength + 1);
OnReceive?.Invoke(cleaned);
break;
}
}
}
catch (Exception ex) {
OnDisconnected?.Invoke(ex.Message);
await ConnectAsync()
.ConfigureAwait(false);
}
}
}
}