Skip to content

Commit 66be5a9

Browse files
mgravellmgravell
committed
push proxy
1 parent 9e51da7 commit 66be5a9

11 files changed

Lines changed: 255 additions & 73 deletions

‎src/RESPite/Streams/ReadStatus.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@ internal enum ReadStatus
88
Faulted,
99
ReadSync,
1010
ReadAsync,
11+
ReadWorker,
1112
TransitioningToAsync,
1213
UpdateWriteTime,
1314
ProcessBuffer,

‎src/RESPite/Streams/RespStream.Read.cs‎

Lines changed: 68 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,6 @@
11
using System.Buffers;
22
using System.Diagnostics;
3+
using System.Net.Sockets;
34
using System.Runtime.CompilerServices;
45
using RESPite.Buffers;
56
using RESPite.Internal;
@@ -26,16 +27,19 @@ private protected virtual void RecordConnectionFailed(StreamFailureKind kind, Ex
2627
public long TotalBytesRead => _totalBytesRead;
2728

2829
[Conditional("PARSE_DETAIL")]
29-
private protected void OnDetailLog(string message) { }
30+
private protected void OnDetailLog(string message)
31+
{
32+
#if PARSE_DETAIL
33+
Console.WriteLine(message);
34+
#endif
35+
}
3036

3137
private async Task ReadAllAsync(Stream tail, CancellationToken cancellationToken)
3238
{
3339
if (_readStatus is not ReadStatus.TransitioningToAsync)
3440
{
3541
// preserve existing state if transitioning
36-
_readStatus = ReadStatus.Init;
37-
_readState = default;
38-
_readBuffer = CycleBuffer.Create(pool: ReaderBufferPool);
42+
InitRead();
3943
}
4044

4145
try
@@ -88,11 +92,70 @@ private async Task ReadAllAsync(Stream tail, CancellationToken cancellationToken
8892
}
8993
}
9094

91-
private void ReadAllSync(Stream tail, CancellationToken cancellationToken)
95+
private protected void InitRead()
9296
{
9397
_readStatus = ReadStatus.Init;
9498
_readState = default;
9599
_readBuffer = CycleBuffer.Create(pool: ReaderBufferPool);
100+
}
101+
102+
private protected Memory<byte> GetReceiveBuffer()
103+
{
104+
_readStatus = ReadStatus.ReadSync;
105+
return _readBuffer.GetUncommittedMemory();
106+
}
107+
108+
private protected bool OnAfterReceive(int read, bool inline)
109+
{
110+
try
111+
{
112+
_readStatus = ReadStatus.UpdateWriteTime;
113+
UpdateLastReadTime();
114+
115+
DebugCounters.OnAsyncRead(read, inline);
116+
_readStatus = ReadStatus.TryParseResult;
117+
return CommitAndParseFrames(read) & !ForceReconnect;
118+
}
119+
catch (Exception ex)
120+
{
121+
OnCleanup(SocketError.Fault, ex);
122+
return false;
123+
}
124+
}
125+
126+
private protected void OnCleanup(SocketError error, Exception? fault = null)
127+
{
128+
try
129+
{
130+
switch (fault)
131+
{
132+
case null when error is SocketError.Success:
133+
_readStatus = ReadStatus.ProcessBufferComplete;
134+
_readBuffer.Release(); // clean exit, we can recycle
135+
136+
_readStatus = ReadStatus.RanToCompletion;
137+
RecordConnectionFailed(StreamFailureKind.SocketClosed);
138+
break;
139+
case null when error is SocketError.ConnectionReset or SocketError.ConnectionAborted:
140+
case EndOfStreamException or OperationCanceledException when _readStatus is ReadStatus.ReadAsync:
141+
_readStatus = ReadStatus.RanToCompletion;
142+
RecordConnectionFailed(StreamFailureKind.SocketClosed);
143+
break;
144+
default:
145+
_readStatus = ReadStatus.Faulted;
146+
RecordConnectionFailed(StreamFailureKind.InternalFailure, fault);
147+
break;
148+
}
149+
}
150+
finally
151+
{
152+
_readBuffer = default; // wipe, however we exited
153+
}
154+
}
155+
156+
private void ReadAllSync(Stream tail, CancellationToken cancellationToken)
157+
{
158+
InitRead();
96159
try
97160
{
98161
int read;

‎toys/RESPite.Proxy/Program.cs‎

Lines changed: 3 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
using System.Buffers;
22
using System.Net;
33
using System.Net.Sockets;
4-
using Microsoft.AspNetCore.Connections;
54
using RESPite.Proxy;
65
using RESPite.Streams;
76

@@ -10,10 +9,12 @@
109
Password = "letmein",
1110
};
1211
using var pool = new WorkerPool();
12+
pool.AddLog(Console.WriteLine);
1313
var proxy = new ProxyServer(proxyOptions, applicationLifetime: null);
1414
using var socket = new ProxySocketServer(proxy, pool);
1515
socket.Start(new IPEndPoint(IPAddress.Loopback, 6380));
16-
16+
Console.WriteLine("Listening...");
17+
Console.ReadLine();
1718
/*
1819
1920
var builder = WebApplication.CreateBuilder(args);

‎toys/RESPite.Proxy/ProxyClient.cs‎

Lines changed: 110 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -3,7 +3,6 @@
33
using System.Diagnostics;
44
using System.IO.Pipelines;
55
using System.Net.Sockets;
6-
using System.Runtime.InteropServices;
76
using RESPite.Buffers;
87
using RESPite.Messages;
98
using RESPite.Streams;
@@ -14,29 +13,38 @@ namespace RESPite.Proxy;
1413
internal sealed class SocketProxyClient : ProxyClient, ICycleBufferCallback
1514
{
1615
private CycleBuffer _buffer;
17-
private StateFlags _stateFlags;
18-
private readonly Socket _outbound;
19-
private readonly WorkerSocketAsyncEventArgs _writeArgs;
16+
private StateFlags _writeFlags;
17+
private readonly Socket _client;
18+
private readonly WorkerSocketAsyncEventArgs _readArgs, _writeArgs;
2019

21-
public SocketProxyClient(ProxyServer.InnerLeg upstream, Socket outbound) : base(upstream)
20+
public SocketProxyClient(ProxyServer.InnerLeg upstream, Socket client) : base(upstream)
2221
{
2322
_buffer = CycleBuffer.Create(pool: upstream.BufferPool, callback: this);
23+
_readArgs = new();
24+
_readArgs.Init(this, WorkerStep.SocketProxyClientReadCallback);
2425
_writeArgs = new();
25-
_outbound = outbound;
26+
_writeArgs.Init(this, WorkerStep.SocketProxyClientWriteCallback);
27+
_client = client;
28+
}
29+
30+
public void StartReading()
31+
{
32+
InitRead();
33+
_readArgs.Pool.Enqueue(this, WorkerStep.SocketProxyClientRead);
2634
}
2735

2836
protected override void SendRawSynchronized(ReadOnlySpan<byte> frame)
2937
{
3038
bool lockTaken = false;
3139
try
3240
{
33-
TakeLock(ref lockTaken);
41+
TakeWriteLock(ref lockTaken);
3442
_buffer.Write(frame);
35-
ActivateInsideLock(StateFlags.Flush);
43+
ActivateWriterInsideLock(StateFlags.Flush);
3644
}
3745
finally
3846
{
39-
ReleaseLock(ref lockTaken);
47+
ReleaseWriteLock(ref lockTaken);
4048
}
4149
}
4250

@@ -49,12 +57,12 @@ private void OnActivate(StateFlags newFlags)
4957
bool lockTaken = false;
5058
try
5159
{
52-
TakeLock(ref lockTaken);
53-
ActivateInsideLock(newFlags);
60+
TakeWriteLock(ref lockTaken);
61+
ActivateWriterInsideLock(newFlags);
5462
}
5563
finally
5664
{
57-
ReleaseLock(ref lockTaken);
65+
ReleaseWriteLock(ref lockTaken);
5866
}
5967
}
6068

@@ -71,6 +79,18 @@ private bool CheckSend()
7179
return false;
7280
}
7381
Debug.Assert(_writeArgs.BytesTransferred == _writeArgs.MemoryBuffer.Length, "incomplete buffer write");
82+
83+
// we're done with the bytes
84+
bool lockTaken = false;
85+
try
86+
{
87+
TakeWriteLock(ref lockTaken);
88+
_buffer.DiscardCommitted(_writeArgs.BytesTransferred);
89+
}
90+
finally
91+
{
92+
ReleaseWriteLock(ref lockTaken);
93+
}
7494
return true;
7595
}
7696

@@ -79,68 +99,123 @@ private void CloseWriter()
7999
bool lockTaken = false;
80100
try
81101
{
82-
TakeLock(ref lockTaken);
83-
_stateFlags = (_stateFlags | StateFlags.Closed) & ~StateFlags.ActiveWriter;
102+
TakeWriteLock(ref lockTaken);
103+
_writeFlags = (_writeFlags | StateFlags.Closed) & ~StateFlags.ActiveWriter;
84104
}
85105
finally
86106
{
87-
ReleaseLock(ref lockTaken);
107+
ReleaseWriteLock(ref lockTaken);
108+
}
109+
}
110+
111+
internal void WorkerRead()
112+
{
113+
try
114+
{
115+
const uint MAX_LOOP = 5;
116+
uint loop = 0; // uint so we're not too concerned about wrap-around
117+
while (true)
118+
{
119+
_readArgs.SetBuffer(GetReceiveBuffer());
120+
if (_client.ReceiveAsync(_readArgs))
121+
return; // gone async, gets reactivated via pool
122+
123+
if (!CheckReceive(inline: true))
124+
break; // validation failed
125+
126+
if (++loop > MAX_LOOP && _readArgs.Pool.HasWork)
127+
{
128+
// yield to pool and come back for more
129+
_readArgs.Pool.Enqueue(this, WorkerStep.SocketProxyClientRead);
130+
break;
131+
}
132+
}
133+
}
134+
catch (Exception ex)
135+
{
136+
OnCleanup(SocketError.Fault, ex);
88137
}
89138
}
90139

140+
private bool CheckReceive(bool inline)
141+
{
142+
var err = _readArgs.SocketError;
143+
if (err is not SocketError.Success)
144+
{
145+
OnCleanup(err);
146+
return false;
147+
}
148+
149+
OnAfterReceive(_readArgs.BytesTransferred, inline);
150+
return true;
151+
}
152+
153+
public void WorkerReadCallback()
154+
{
155+
if (CheckReceive(inline: false)) WorkerRead(); // try to do more
156+
}
157+
91158
internal void WorkerWrite()
92159
{
93160
bool lockTaken = false;
94161
try
95162
{
96-
const int maxLoop = 5;
97-
int loop = 0;
98-
do
163+
const uint MAX_LOOP = 5;
164+
uint loop = 0;
165+
while (true)
99166
{
100-
TakeLock(ref lockTaken);
101-
var minBytes = (_stateFlags & StateFlags.Flush) == 0 ? -1 : 1;
167+
TakeWriteLock(ref lockTaken);
168+
var minBytes = (_writeFlags & StateFlags.Flush) == 0 ? -1 : 1;
102169
var success = _buffer.TryGetFirstCommittedMemory(minBytes, out var memory);
103170
if (!success)
104171
{
105-
_stateFlags &= ~StateFlags.ActiveWriter;
172+
_writeFlags &= ~StateFlags.ActiveWriter;
106173
break;
107174
}
108175

109-
ReleaseLock(ref lockTaken);
176+
ReleaseWriteLock(ref lockTaken);
110177

111-
_writeArgs.Step = WorkerStep.SocketProxyClientWriteCallback;
112-
_writeArgs.SetBuffer(MemoryMarshal.AsMemory(memory));
113-
if (_outbound.SendAsync(_writeArgs)) break; // gone async, gets reactivated via pool
178+
_writeArgs.SetBuffer(memory);
179+
if (_client.SendAsync(_writeArgs))
180+
break; // gone async, gets reactivated via pool
181+
182+
if (!CheckSend())
183+
break; // validation failed
184+
185+
if (++loop > MAX_LOOP && _writeArgs.Pool.HasWork)
186+
{
187+
// yield to pool and come back for more
188+
_writeArgs.Pool.Enqueue(this, WorkerStep.SocketProxyClientWrite);
189+
break;
190+
}
114191
}
115-
// on sync write, assert and continue - but don't hog the pool forever
116-
while (CheckSend() & ++loop <= maxLoop);
117192
}
118193
finally
119194
{
120-
ReleaseLock(ref lockTaken);
195+
ReleaseWriteLock(ref lockTaken);
121196
}
122197
}
123198

124-
private void ActivateInsideLock(StateFlags newFlags)
199+
private void ActivateWriterInsideLock(StateFlags newFlags)
125200
{
126-
Debug.Assert(Monitor.IsEntered(this), $"{nameof(ActivateInsideLock)} must be called while holding the writer lock.");
201+
Debug.Assert(Monitor.IsEntered(this), $"{nameof(ActivateWriterInsideLock)} must be called while holding the writer lock.");
127202

128-
var state = _stateFlags;
203+
var state = _writeFlags;
129204
if ((state & StateFlags.Closed) != 0) return;
130205
state |= newFlags & ~StateFlags.ActiveWriter;
131206
if ((state & StateFlags.ActiveWriter) == 0)
132207
{
133208
state |= StateFlags.ActiveWriter;
134-
_stateFlags = state;
209+
_writeFlags = state;
135210
_writeArgs.Pool.Enqueue(this, WorkerStep.SocketProxyClientWrite);
136211
}
137212
else
138213
{
139-
_stateFlags = state;
214+
_writeFlags = state;
140215
}
141216
}
142217

143-
private void TakeLock(ref bool lockTaken)
218+
private void TakeWriteLock(ref bool lockTaken)
144219
{
145220
if (!lockTaken)
146221
{
@@ -150,7 +225,7 @@ private void TakeLock(ref bool lockTaken)
150225
static void Throw() => throw new TimeoutException("Unable to acquire writer lock");
151226
}
152227

153-
private void ReleaseLock(ref bool lockTaken)
228+
private void ReleaseWriteLock(ref bool lockTaken)
154229
{
155230
if (lockTaken)
156231
{

‎toys/RESPite.Proxy/ProxyServer.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -162,7 +162,7 @@ public void RunClient(Socket socket)
162162
WorkerPool.DebugAssertWorker();
163163
var client = new SocketProxyClient(this, socket);
164164
server.RegisterClient(client);
165-
return client.ExecuteAsync();
165+
client.StartReading();
166166
}
167167
}
168168

‎toys/RESPite.Proxy/SocketClient.cs‎

Lines changed: 0 additions & 9 deletions
This file was deleted.

0 commit comments

Comments
 (0)