scaling

mgravell · mgravell · commit a4b9c7c8c9c9 · 2026-07-03T13:16:38.000+01:00
diff --git a/src/RESPite/Streams/RespStream.Read.cs b/src/RESPite/Streams/RespStream.Read.cs
@@ -118,12 +118,12 @@ private protected bool OnAfterReceive(int read, bool inline)
         }
         catch (Exception ex)
         {
-            OnCleanup(SocketError.Fault, ex);
+            OnReceiveCleanup(SocketError.Fault, ex);
             return false;
         }
     }
 
-    private protected void OnCleanup(SocketError error, Exception? fault = null)
+    private protected virtual void OnReceiveCleanup(SocketError error, Exception? fault = null)
     {
         try
         {
diff --git a/toys/RESPite.Proxy/Program.cs b/toys/RESPite.Proxy/Program.cs
@@ -7,14 +7,29 @@
 var proxyOptions = new ProxyServerOptions
 {
     Password = "letmein",
+    UpstreamConnectionCount = 5,
 };
-using var pool = new WorkerPool();
-pool.AddLog(Console.WriteLine);
+using var pool = new WorkerPool(workers: 0); // use CPU count
+pool.AddDebugLog(Console.WriteLine);
 var proxy = new ProxyServer(proxyOptions, applicationLifetime: null);
 using var socket = new ProxySocketServer(proxy, pool);
 socket.Start(new IPEndPoint(IPAddress.Loopback, 6380));
-Console.WriteLine("Listening...");
-Console.ReadLine();
+
+ulong lastOpCount = 0;
+int lastActiveClients = 0;
+bool first = true;
+while (true)
+{
+    var opCount = proxy.GetOpCount(out var activeClients);
+    if (first || opCount != lastOpCount || activeClients != lastActiveClients)
+    {
+        Console.WriteLine($"Active clients: {activeClients}; commands processed: {opCount}");
+        lastOpCount = opCount;
+        lastActiveClients = activeClients;
+        first = false;
+    }
+    await Task.Delay(TimeSpan.FromSeconds(5));
+}
 /*
 
 var builder = WebApplication.CreateBuilder(args);
diff --git a/toys/RESPite.Proxy/ProxyClient.cs b/toys/RESPite.Proxy/ProxyClient.cs
@@ -12,14 +12,14 @@ namespace RESPite.Proxy;
 
 internal sealed class SocketProxyClient : ProxyClient, ICycleBufferCallback
 {
-    private CycleBuffer _buffer;
+    private CycleBuffer _receiveBuffer;
     private StateFlags _writeFlags;
     private readonly Socket _client;
     private readonly WorkerSocketAsyncEventArgs _readArgs, _writeArgs;
 
     public SocketProxyClient(ProxyServer.InnerLeg upstream, Socket client) : base(upstream)
     {
-        _buffer = CycleBuffer.Create(pool: upstream.BufferPool, callback: this);
+        _receiveBuffer = CycleBuffer.Create(pool: upstream.BufferPool, callback: this);
         _readArgs = new();
         _readArgs.Init(this, WorkerStep.SocketProxyClientReadCallback);
         _writeArgs = new();
@@ -39,7 +39,8 @@ protected override void SendRawSynchronized(ReadOnlySpan<byte> frame)
         try
         {
             TakeWriteLock(ref lockTaken);
-            _buffer.Write(frame);
+            if ((_writeFlags & StateFlags.Closed) != 0) return; // torn down; nothing left to write to
+            _receiveBuffer.Write(frame);
             ActivateWriterInsideLock(StateFlags.Flush);
         }
         finally
@@ -78,14 +79,17 @@ private bool CheckSend()
             CloseWriter();
             return false;
         }
-        Debug.Assert(_writeArgs.BytesTransferred == _writeArgs.MemoryBuffer.Length, "incomplete buffer write");
+        // a stream socket may legally report a partial send; we only discard what actually went out,
+        // and the remainder stays committed for the next TryGetFirstCommittedMemory, so this is safe.
+        Debug.Assert(_writeArgs.BytesTransferred <= _writeArgs.MemoryBuffer.Length, "over-send?!");
 
-        // we're done with the bytes
+        // we're done with the bytes that made it onto the wire
         bool lockTaken = false;
         try
         {
             TakeWriteLock(ref lockTaken);
-            _buffer.DiscardCommitted(_writeArgs.BytesTransferred);
+            if ((_writeFlags & StateFlags.Closed) != 0) return false; // torn down; buffer already released
+            _receiveBuffer.DiscardCommitted(_writeArgs.BytesTransferred);
         }
         finally
         {
@@ -106,6 +110,15 @@ private void CloseWriter()
         {
             ReleaseWriteLock(ref lockTaken);
         }
+        _writeArgs.Dispose();
+    }
+
+    protected override void ReleaseResources()
+    {
+        // dispose the socket first: this aborts any in-flight send/receive so the SAEAs are no longer
+        // in use by the time we dispose them (a stray completion just lands as OperationAborted).
+        try { _client.Dispose(); }
+        catch { /* already gone */ }
     }
 
     internal void WorkerRead()
@@ -131,23 +144,41 @@ internal void WorkerRead()
                 }
             }
         }
+        catch (Exception ex) when (ex is ObjectDisposedException or OperationCanceledException)
+        {
+            // socket torn down during shutdown/cancellation: treat as a normal close, not a fault
+            OnReceiveCleanup(SocketError.ConnectionAborted);
+        }
         catch (Exception ex)
         {
-            OnCleanup(SocketError.Fault, ex);
+            OnReceiveCleanup(SocketError.Fault, ex);
         }
     }
 
+    private protected override void OnReceiveCleanup(SocketError error, Exception? fault = null)
+    {
+        base.OnReceiveCleanup(error, fault);
+        _readArgs.Dispose();
+        _receiveBuffer.Release();
+    }
+
     private bool CheckReceive(bool inline)
     {
         var err = _readArgs.SocketError;
         if (err is not SocketError.Success)
         {
-            OnCleanup(err);
+            OnReceiveCleanup(err);
             return false;
         }
 
-        OnAfterReceive(_readArgs.BytesTransferred, inline);
-        return true;
+        if (_readArgs.BytesTransferred == 0)
+        {
+            // EOF
+            OnReceiveCleanup(SocketError.Success);
+            return false;
+        }
+
+        return OnAfterReceive(_readArgs.BytesTransferred, inline);
     }
 
     public void WorkerReadCallback()
@@ -165,8 +196,13 @@ internal void WorkerWrite()
             while (true)
             {
                 TakeWriteLock(ref lockTaken);
+                if ((_writeFlags & StateFlags.Closed) != 0)
+                {
+                    _writeFlags &= ~StateFlags.ActiveWriter;
+                    break; // torn down (buffer released); stop pumping
+                }
                 var minBytes = (_writeFlags & StateFlags.Flush) == 0 ? -1 : 1;
-                var success = _buffer.TryGetFirstCommittedMemory(minBytes, out var memory);
+                var success = _receiveBuffer.TryGetFirstCommittedMemory(minBytes, out var memory);
                 if (!success)
                 {
                     _writeFlags &= ~StateFlags.ActiveWriter;
@@ -198,7 +234,7 @@ internal void WorkerWrite()
 
     private void ActivateWriterInsideLock(StateFlags newFlags)
     {
-        Debug.Assert(Monitor.IsEntered(this), $"{nameof(ActivateWriterInsideLock)} must be called while holding the writer lock.");
+        Debug.Assert(Monitor.IsEntered(_writeLock), $"{nameof(ActivateWriterInsideLock)} must be called while holding the writer lock.");
 
         var state = _writeFlags;
         if ((state & StateFlags.Closed) != 0) return;
@@ -215,11 +251,16 @@ private void ActivateWriterInsideLock(StateFlags newFlags)
         }
     }
 
+    // guards _buffer/_writeFlags on the write side; a dedicated object (rather than 'this') so no
+    // external code can contend the lock, and reentrancy via ICycleBufferCallback.PageComplete is
+    // still safe (Monitor is reentrant on the same thread).
+    private readonly object _writeLock = new();
+
     private void TakeWriteLock(ref bool lockTaken)
     {
         if (!lockTaken)
         {
-            Monitor.TryEnter(this, 10_000,  ref lockTaken);
+            Monitor.TryEnter(_writeLock, 10_000, ref lockTaken);
             if (!lockTaken) Throw();
         }
         static void Throw() => throw new TimeoutException("Unable to acquire writer lock");
@@ -229,7 +270,7 @@ private void ReleaseWriteLock(ref bool lockTaken)
     {
         if (lockTaken)
         {
-            Monitor.Exit(this);
+            Monitor.Exit(_writeLock);
             lockTaken = false;
         }
     }
@@ -259,9 +300,12 @@ internal abstract class ProxyClient(ProxyServer.InnerLeg upstream) : RespStream
     protected CancellationToken Lifetime => upstream.Lifetime;
     public int Id { get; set; }
     public int Database => _db;
+    public ulong OpCount => Interlocked.Read(ref _opCount);
+
     private int _db;
     private TaskCompletionSource _completionSource = new(TaskCreationOptions.RunContinuationsAsynchronously);
 
+    private ulong _opCount;
     private bool OnSelect(int db)
     {
         if (db < 0 | db > 999_999_999) return false;
@@ -275,8 +319,14 @@ public Task ExecuteAsync(PipeReader source)
         return _completionSource.Task;
     }
 
+    private int _closed; // 0 = open, 1 = closed; guards teardown so it runs exactly once
+
     private protected override void RecordConnectionFailed(StreamFailureKind kind, Exception? fault = null)
     {
+        // teardown can be reached from the read pump *and*, once the worker is multi-threaded, from a
+        // write-side fault on another thread; make it idempotent so we don't double-remove/double-dispose
+        if (Interlocked.Exchange(ref _closed, 1) != 0) return;
+
         if (fault is null)
         {
             _completionSource.TrySetResult();
@@ -287,13 +337,22 @@ private protected override void RecordConnectionFailed(StreamFailureKind kind, E
         }
 
         upstream.Remove(this);
+        ReleaseResources();
     }
 
+    /// <summary>
+    /// Releases connection-scoped resources (sockets, buffers, ...). Invoked exactly once, after the
+    /// connection has been removed from the pool.
+    /// </summary>
+    protected virtual void ReleaseResources() { }
+
     protected override unsafe void OnReadFrame(
         RespPrefix prefix,
         ReadOnlySpan<byte> frame,
         ref IMemoryOwner<byte>? memoryOwner)
     {
+        _opCount++;
+
         ReadOnlyMemory<byte> localResponse = default;
         IDisposable? lease = null;
         KnownCommands command = KnownCommands.Unknown;
diff --git a/toys/RESPite.Proxy/ProxyServer.cs b/toys/RESPite.Proxy/ProxyServer.cs
@@ -71,7 +71,29 @@ internal void RegisterClient(ProxyClient client)
 
     internal bool TryGetClient(int id, out ProxyClient client) => _clients.TryGetValue(id, out client!);
 
-    internal void RemoveClient(ProxyClient client) => _clients.TryRemove(client.Id, out _);
+    internal void RemoveClient(ProxyClient client)
+    {
+        var opCount = client.OpCount;
+        if (_clients.TryRemove(client.Id, out _))
+        {
+            Interlocked.Add(ref _opCountFromDeadConnections, opCount);
+        }
+    }
+
+    public ulong GetOpCount(out int activeClients)
+    {
+        ulong tally = 0;
+        activeClients = 0;
+        foreach (var client in _clients)
+        {
+            tally += client.Value.OpCount;
+            activeClients++;
+        }
+
+        // intentionally delay this read
+        return tally + Interlocked.Read(ref _opCountFromDeadConnections);
+    }
+    private ulong _opCountFromDeadConnections;
 
     internal sealed class InnerLeg(ProxyServer server, Stream tail) : RespStream
     {
diff --git a/toys/RESPite.Proxy/WorkerPool.cs b/toys/RESPite.Proxy/WorkerPool.cs

Original file line number	Diff line number	Diff line change
`@@ -118,12 +118,12 @@ private protected bool OnAfterReceive(int read, bool inline)`
`118`	`118`	`}`
`119`	`119`	`catch (Exception ex)`
`120`	`120`	`{`
`121`		`- OnCleanup(SocketError.Fault, ex);`
	`121`	`+ OnReceiveCleanup(SocketError.Fault, ex);`
`122`	`122`	`return false;`
`123`	`123`	`}`
`124`	`124`	`}`
`125`	`125`
`126`		`- private protected void OnCleanup(SocketError error, Exception? fault = null)`
	`126`	`+ private protected virtual void OnReceiveCleanup(SocketError error, Exception? fault = null)`
`127`	`127`	`{`
`128`	`128`	`try`
`129`	`129`	`{`
Original file line number	Diff line number	Diff line change
`@@ -12,14 +12,14 @@ namespace RESPite.Proxy;`
`12`	`12`
`13`	`13`	`internal sealed class SocketProxyClient : ProxyClient, ICycleBufferCallback`
`14`	`14`	`{`
`15`		`- private CycleBuffer _buffer;`
	`15`	`+ private CycleBuffer _receiveBuffer;`
`16`	`16`	`private StateFlags _writeFlags;`
`17`	`17`	`private readonly Socket _client;`
`18`	`18`	`private readonly WorkerSocketAsyncEventArgs _readArgs, _writeArgs;`
`19`	`19`
`20`	`20`	`public SocketProxyClient(ProxyServer.InnerLeg upstream, Socket client) : base(upstream)`
`21`	`21`	`{`
`22`		`- _buffer = CycleBuffer.Create(pool: upstream.BufferPool, callback: this);`
	`22`	`+ _receiveBuffer = CycleBuffer.Create(pool: upstream.BufferPool, callback: this);`
`23`	`23`	`_readArgs = new();`
`24`	`24`	`_readArgs.Init(this, WorkerStep.SocketProxyClientReadCallback);`
`25`	`25`	`_writeArgs = new();`
`@@ -39,7 +39,8 @@ protected override void SendRawSynchronized(ReadOnlySpan<byte> frame)`
`39`	`39`	`try`
`40`	`40`	`{`
`41`	`41`	`TakeWriteLock(ref lockTaken);`
`42`		`- _buffer.Write(frame);`
	`42`	`+ if ((_writeFlags & StateFlags.Closed) != 0) return; // torn down; nothing left to write to`
	`43`	`+ _receiveBuffer.Write(frame);`
`43`	`44`	`ActivateWriterInsideLock(StateFlags.Flush);`
`44`	`45`	`}`
`45`	`46`	`finally`
`@@ -78,14 +79,17 @@ private bool CheckSend()`
`78`	`79`	`CloseWriter();`
`79`	`80`	`return false;`
`80`	`81`	`}`
`81`		`- Debug.Assert(_writeArgs.BytesTransferred == _writeArgs.MemoryBuffer.Length, "incomplete buffer write");`
	`82`	`+ // a stream socket may legally report a partial send; we only discard what actually went out,`
	`83`	`+ // and the remainder stays committed for the next TryGetFirstCommittedMemory, so this is safe.`
	`84`	`+ Debug.Assert(_writeArgs.BytesTransferred <= _writeArgs.MemoryBuffer.Length, "over-send?!");`
`82`	`85`
`83`		`- // we're done with the bytes`
	`86`	`+ // we're done with the bytes that made it onto the wire`
`84`	`87`	`bool lockTaken = false;`
`85`	`88`	`try`
`86`	`89`	`{`
`87`	`90`	`TakeWriteLock(ref lockTaken);`
`88`		`- _buffer.DiscardCommitted(_writeArgs.BytesTransferred);`
	`91`	`+ if ((_writeFlags & StateFlags.Closed) != 0) return false; // torn down; buffer already released`
	`92`	`+ _receiveBuffer.DiscardCommitted(_writeArgs.BytesTransferred);`
`89`	`93`	`}`
`90`	`94`	`finally`
`91`	`95`	`{`
`@@ -106,6 +110,15 @@ private void CloseWriter()`
`106`	`110`	`{`
`107`	`111`	`ReleaseWriteLock(ref lockTaken);`
`108`	`112`	`}`
	`113`	`+ _writeArgs.Dispose();`
	`114`	`+ }`
	`115`	`+`
	`116`	`+ protected override void ReleaseResources()`
	`117`	`+ {`
	`118`	`+ // dispose the socket first: this aborts any in-flight send/receive so the SAEAs are no longer`
	`119`	`+ // in use by the time we dispose them (a stray completion just lands as OperationAborted).`
	`120`	`+ try { _client.Dispose(); }`
	`121`	`+ catch { /* already gone */ }`
`109`	`122`	`}`
`110`	`123`
`111`	`124`	`internal void WorkerRead()`
`@@ -131,23 +144,41 @@ internal void WorkerRead()`
`131`	`144`	`}`
`132`	`145`	`}`
`133`	`146`	`}`
	`147`	`+ catch (Exception ex) when (ex is ObjectDisposedException or OperationCanceledException)`
	`148`	`+ {`
	`149`	`+ // socket torn down during shutdown/cancellation: treat as a normal close, not a fault`
	`150`	`+ OnReceiveCleanup(SocketError.ConnectionAborted);`
	`151`	`+ }`
`134`	`152`	`catch (Exception ex)`
`135`	`153`	`{`
`136`		`- OnCleanup(SocketError.Fault, ex);`
	`154`	`+ OnReceiveCleanup(SocketError.Fault, ex);`
`137`	`155`	`}`
`138`	`156`	`}`
`139`	`157`
	`158`	`+ private protected override void OnReceiveCleanup(SocketError error, Exception? fault = null)`
	`159`	`+ {`
	`160`	`+ base.OnReceiveCleanup(error, fault);`
	`161`	`+ _readArgs.Dispose();`
	`162`	`+ _receiveBuffer.Release();`
	`163`	`+ }`
	`164`	`+`
`140`	`165`	`private bool CheckReceive(bool inline)`
`141`	`166`	`{`
`142`	`167`	`var err = _readArgs.SocketError;`
`143`	`168`	`if (err is not SocketError.Success)`
`144`	`169`	`{`
`145`		`- OnCleanup(err);`
	`170`	`+ OnReceiveCleanup(err);`
`146`	`171`	`return false;`
`147`	`172`	`}`
`148`	`173`
`149`		`- OnAfterReceive(_readArgs.BytesTransferred, inline);`
`150`		`- return true;`
	`174`	`+ if (_readArgs.BytesTransferred == 0)`
	`175`	`+ {`
	`176`	`+ // EOF`
	`177`	`+ OnReceiveCleanup(SocketError.Success);`
	`178`	`+ return false;`
	`179`	`+ }`
	`180`	`+`
	`181`	`+ return OnAfterReceive(_readArgs.BytesTransferred, inline);`
`151`	`182`	`}`
`152`	`183`
`153`	`184`	`public void WorkerReadCallback()`
`@@ -165,8 +196,13 @@ internal void WorkerWrite()`
`165`	`196`	`while (true)`
`166`	`197`	`{`
`167`	`198`	`TakeWriteLock(ref lockTaken);`
	`199`	`+ if ((_writeFlags & StateFlags.Closed) != 0)`
	`200`	`+ {`
	`201`	`+ _writeFlags &= ~StateFlags.ActiveWriter;`
	`202`	`+ break; // torn down (buffer released); stop pumping`
	`203`	`+ }`
`168`	`204`	`var minBytes = (_writeFlags & StateFlags.Flush) == 0 ? -1 : 1;`
`169`		`- var success = _buffer.TryGetFirstCommittedMemory(minBytes, out var memory);`
	`205`	`+ var success = _receiveBuffer.TryGetFirstCommittedMemory(minBytes, out var memory);`
`170`	`206`	`if (!success)`
`171`	`207`	`{`
`172`	`208`	`_writeFlags &= ~StateFlags.ActiveWriter;`
`@@ -198,7 +234,7 @@ internal void WorkerWrite()`
`198`	`234`
`199`	`235`	`private void ActivateWriterInsideLock(StateFlags newFlags)`
`200`	`236`	`{`
`201`		`- Debug.Assert(Monitor.IsEntered(this), $"{nameof(ActivateWriterInsideLock)} must be called while holding the writer lock.");`
	`237`	`+ Debug.Assert(Monitor.IsEntered(_writeLock), $"{nameof(ActivateWriterInsideLock)} must be called while holding the writer lock.");`
`202`	`238`
`203`	`239`	`var state = _writeFlags;`
`204`	`240`	`if ((state & StateFlags.Closed) != 0) return;`
`@@ -215,11 +251,16 @@ private void ActivateWriterInsideLock(StateFlags newFlags)`
`215`	`251`	`}`
`216`	`252`	`}`
`217`	`253`
	`254`	`+ // guards _buffer/_writeFlags on the write side; a dedicated object (rather than 'this') so no`
	`255`	`+ // external code can contend the lock, and reentrancy via ICycleBufferCallback.PageComplete is`
	`256`	`+ // still safe (Monitor is reentrant on the same thread).`
	`257`	`+ private readonly object _writeLock = new();`
	`258`	`+`
`218`	`259`	`private void TakeWriteLock(ref bool lockTaken)`
`219`	`260`	`{`
`220`	`261`	`if (!lockTaken)`
`221`	`262`	`{`
`222`		`- Monitor.TryEnter(this, 10_000, ref lockTaken);`
	`263`	`+ Monitor.TryEnter(_writeLock, 10_000, ref lockTaken);`
`223`	`264`	`if (!lockTaken) Throw();`
`224`	`265`	`}`
`225`	`266`	`static void Throw() => throw new TimeoutException("Unable to acquire writer lock");`
`@@ -229,7 +270,7 @@ private void ReleaseWriteLock(ref bool lockTaken)`
`229`	`270`	`{`
`230`	`271`	`if (lockTaken)`
`231`	`272`	`{`
`232`		`- Monitor.Exit(this);`
	`273`	`+ Monitor.Exit(_writeLock);`
`233`	`274`	`lockTaken = false;`
`234`	`275`	`}`
`235`	`276`	`}`
`@@ -259,9 +300,12 @@ internal abstract class ProxyClient(ProxyServer.InnerLeg upstream) : RespStream`
`259`	`300`	`protected CancellationToken Lifetime => upstream.Lifetime;`
`260`	`301`	`public int Id { get; set; }`
`261`	`302`	`public int Database => _db;`
	`303`	`+ public ulong OpCount => Interlocked.Read(ref _opCount);`
	`304`	`+`
`262`	`305`	`private int _db;`
`263`	`306`	`private TaskCompletionSource _completionSource = new(TaskCreationOptions.RunContinuationsAsynchronously);`
`264`	`307`
	`308`	`+ private ulong _opCount;`
`265`	`309`	`private bool OnSelect(int db)`
`266`	`310`	`{`
`267`	`311`	`if (db < 0 \| db > 999_999_999) return false;`
`@@ -275,8 +319,14 @@ public Task ExecuteAsync(PipeReader source)`
`275`	`319`	`return _completionSource.Task;`
`276`	`320`	`}`
`277`	`321`
	`322`	`+ private int _closed; // 0 = open, 1 = closed; guards teardown so it runs exactly once`
	`323`	`+`
`278`	`324`	`private protected override void RecordConnectionFailed(StreamFailureKind kind, Exception? fault = null)`
`279`	`325`	`{`
	`326`	`+ // teardown can be reached from the read pump and, once the worker is multi-threaded, from a`
	`327`	`+ // write-side fault on another thread; make it idempotent so we don't double-remove/double-dispose`
	`328`	`+ if (Interlocked.Exchange(ref _closed, 1) != 0) return;`
	`329`	`+`
`280`	`330`	`if (fault is null)`
`281`	`331`	`{`
`282`	`332`	`_completionSource.TrySetResult();`
`@@ -287,13 +337,22 @@ private protected override void RecordConnectionFailed(StreamFailureKind kind, E`
`287`	`337`	`}`
`288`	`338`
`289`	`339`	`upstream.Remove(this);`
	`340`	`+ ReleaseResources();`
`290`	`341`	`}`
`291`	`342`
	`343`	`+ /// <summary>`
	`344`	`+ /// Releases connection-scoped resources (sockets, buffers, ...). Invoked exactly once, after the`
	`345`	`+ /// connection has been removed from the pool.`
	`346`	`+ /// </summary>`
	`347`	`+ protected virtual void ReleaseResources() { }`
	`348`	`+`
`292`	`349`	`protected override unsafe void OnReadFrame(`
`293`	`350`	`RespPrefix prefix,`
`294`	`351`	`ReadOnlySpan<byte> frame,`
`295`	`352`	`ref IMemoryOwner<byte>? memoryOwner)`
`296`	`353`	`{`
	`354`	`+ _opCount++;`
	`355`	`+`
`297`	`356`	`ReadOnlyMemory<byte> localResponse = default;`
`298`	`357`	`IDisposable? lease = null;`
`299`	`358`	`KnownCommands command = KnownCommands.Unknown;`