tidy up kestrel code; flush more eagerly when writing

KestrelRedisServer/RedisConnectionHandler.cs

 using System.Threading.Tasks;
 using Microsoft.AspNetCore.Connections;
-using Microsoft.Extensions.Logging;
 using StackExchange.Redis.Server;
 
 namespace KestrelRedisServer
 {
     public class RedisConnectionHandler : ConnectionHandler
     {
-        private readonly MemoryCacheRedisServer _server;
-        public RedisConnectionHandler(ILogger<RedisConnectionHandler> logger)
-        {
-            _server = new MemoryCacheRedisServer();
-        }
-        public override async Task OnConnectedAsync(ConnectionContext connection)
-        {
-            var client = _server.AddClient();
-            try
-            {
-                while (true)
-                {
-                    var read = await connection.Transport.Input.ReadAsync();
-                    var buffer = read.Buffer;
-                    bool makingProgress = false;
-                    while (_server.TryProcessRequest(ref buffer, client, connection.Transport.Output))
-                    {
-                        makingProgress = true;
-                        await connection.Transport.Output.FlushAsync();
-                    }
-                    connection.Transport.Input.AdvanceTo(buffer.Start, buffer.End);
-
-                    if (!makingProgress && read.IsCompleted) break;
-                }
-            }
-            catch (ConnectionResetException) { } // swallow
-            finally
-            {
-                _server.RemoveClient(client);
-                connection.Transport.Input.Complete();
-                connection.Transport.Output.Complete();
-            }
-        }
+        private readonly RespServer _server;
+        public RedisConnectionHandler(RespServer server) => _server = server;
+        public override Task OnConnectedAsync(ConnectionContext connection)
+            => _server.RunClient(connection.Transport);
     }
 }

KestrelRedisServer/Startup.cs

-using Microsoft.AspNetCore.Builder;
+using System;
+using Microsoft.AspNetCore.Builder;
 using Microsoft.AspNetCore.Hosting;
 using Microsoft.AspNetCore.Http;
 using Microsoft.Extensions.DependencyInjection;
+using StackExchange.Redis.Server;
 
 namespace KestrelRedisServer
 {
-    public class Startup
+    public class Startup : IDisposable
     {
+        RespServer _server = new MemoryCacheRedisServer();
+
         // This method gets called by the runtime. Use this method to add services to the container.
         // For more information on how to configure your application, visit https://go.microsoft.com/fwlink/?LinkID=398940
         public void ConfigureServices(IServiceCollection services)
-        {
-        }
+            => services.Add(new ServiceDescriptor(typeof(RespServer), _server));
+
+        public void Dispose() => _server.Dispose();
 
         // This method gets called by the runtime. Use this method to configure the HTTP request pipeline.
         public void Configure(IApplicationBuilder app, IHostingEnvironment env)
         {
-            if (env.IsDevelopment())
-            {
-                app.UseDeveloperExceptionPage();
-            }
-
-            app.Run(async (context) =>
-            {
-                await context.Response.WriteAsync("Redis-ish server should be running");
-            });
+            if (env.IsDevelopment()) app.UseDeveloperExceptionPage();
+            app.Run(context => context.Response.WriteAsync(_server.GetStats()));
         }
     }
 }

StackExchange.Redis.Server/RedisServer.cs

@@ -16,12 +16,29 @@ public static bool IsMatch(string pattern, string key)
         protected RedisServer(int databases = 16, TextWriter output = null) : base(output)
         {
             if (databases < 1) throw new ArgumentOutOfRangeException(nameof(databases));
+            Databases = databases;
             var config = ServerConfiguration;
             config["timeout"] = "0";
             config["slave-read-only"] = "yes";
             config["databases"] = databases.ToString();
             config["slaveof"] = "";
         }
+        protected override void AppendStats(StringBuilder sb)
+        {
+            base.AppendStats(sb);
+            sb.Append("Databases: ").Append(Databases).AppendLine();
+            lock (ServerSyncLock)
+            {
+                for (int i = 0; i < Databases; i++)
+                {
+                    try
+                    {
+                        sb.Append("Database ").Append(i).Append(": ").Append(Dbsize(i)).AppendLine(" keys");
+                    }
+                    catch { }
+                }
+            }
+        }
         public int Databases { get; }
 
         [RedisCommand(-3)]
@@ -374,7 +391,7 @@ StringBuilder AddHeader()
                     break;
                 case "Stats":
                     AddHeader().Append("total_connections_received:").Append(TotalClientCount).AppendLine()
-                        .Append("total_commands_processed:").Append(CommandsProcesed).AppendLine();
+                        .Append("total_commands_processed:").Append(TotalCommandsProcesed).AppendLine();
                     break;
                 case "Replication":
                     AddHeader().AppendLine("role:master");

StackExchange.Redis.Server/RespServer.cs

@@ -59,6 +59,21 @@ RedisCommandAttribute CheckSignatureAndGetAttribute(MethodInfo method)
             }
             return result;
         }
+
+        public string GetStats()
+        {
+            var sb = new StringBuilder();
+            AppendStats(sb);
+            return sb.ToString();
+        }
+        protected virtual void AppendStats(StringBuilder sb)
+        {
+            sb.Append("Current clients:\t").Append(ClientCount).AppendLine()
+                .Append("Total clients:\t").Append(TotalClientCount).AppendLine()
+                .Append("Total operations:\t").Append(TotalCommandsProcesed).AppendLine()
+                .Append("Error replies:\t").Append(TotalErrorCount).AppendLine();
+        }
+
         [AttributeUsage(AttributeTargets.Method, AllowMultiple = false, Inherited = true)]
         protected sealed class RedisCommandAttribute : Attribute
         {
@@ -171,8 +186,9 @@ protected int TcpPort()
         }
 
         private Action<object> _runClientCallback;
+        // KeepAlive here just to make the compiler happy that we've done *something* with the task
         private Action<object> RunClientCallback => _runClientCallback ??
-            (_runClientCallback = state => RunClient((RedisClient)state));
+            (_runClientCallback = state => GC.KeepAlive(RunClient((IDuplexPipe)state)));
 
         public void Listen(
             EndPoint endpoint,
@@ -212,8 +228,8 @@ public RedisClient AddClient()
             var client = CreateClient();
             lock (_clients)
             {
-                client.Id = ++_nextId;
                 ThrowIfShutdown();
+                client.Id = ++_nextId;
                 _clients.Add(client);
                 TotalClientCount++;
             }
@@ -221,7 +237,7 @@ public RedisClient AddClient()
         }
         public bool RemoveClient(RedisClient client)
         {
-            if (client == null) throw new ArgumentNullException(nameof(client));
+            if (client == null) return false;
             lock (_clients)
             {
                 client.Closed = true;
@@ -237,16 +253,14 @@ private async void ListenForConnections(PipeOptions sendOptions, PipeOptions rec
                     var client = await _listener.AcceptAsync();
                     SocketConnection.SetRecommendedServerOptions(client);
                     var pipe = SocketConnection.Create(client, sendOptions, receiveOptions);
-                    var c = AddClient();
-                    c.LinkedPipe = pipe;
-                    StartOnScheduler(receiveOptions.ReaderScheduler, RunClientCallback, c);
+                    StartOnScheduler(receiveOptions.ReaderScheduler, RunClientCallback, pipe);
                 }
             }
             catch (NullReferenceException) { }
             catch (ObjectDisposedException) { }
             catch (Exception ex)
             {
-                if(!_isShutdown) Log("Listener faulted: " + ex.Message);
+                if (!_isShutdown) Log("Listener faulted: " + ex.Message);
             }
         }
 
@@ -281,33 +295,27 @@ protected virtual void Dispose(bool disposing)
             }
         }
 
-        async void RunClient(RedisClient client)
+        public async Task RunClient(IDuplexPipe pipe)
         {
-            ThrowIfShutdown();
-
-            var input = client?.LinkedPipe?.Input;
-            var output = client?.LinkedPipe?.Output;
-            if (input == null || output == null) return; // nope
-
-
             Exception fault = null;
+            RedisClient client = null;
             try
             {
+                client = AddClient();
                 while (!client.Closed)
                 {
-                    var readResult = await input.ReadAsync();
+                    var readResult = await pipe.Input.ReadAsync();
                     var buffer = readResult.Buffer;
 
                     bool makingProgress = false;
-                    while (!client.Closed && TryProcessRequest(ref buffer, client, output))
+                    while (!client.Closed && await TryProcessRequestAsync(ref buffer, client, pipe.Output))
                     {
                         makingProgress = true;
-                        await output.FlushAsync();
                     }
-                    input.AdvanceTo(buffer.Start, buffer.End);
+                    pipe.Input.AdvanceTo(buffer.Start, buffer.End);
 
                     if (!makingProgress && readResult.IsCompleted)
-                    {
+                    { // nothing to do, and nothing more will be arriving
                         break;
                     }
                 }
@@ -317,8 +325,9 @@ async void RunClient(RedisClient client)
             catch (Exception ex) { fault = ex; }
             finally
             {
-                try { input.Complete(fault); } catch { }
-                try { output.Complete(fault); } catch { }
+                RemoveClient(client);
+                try { pipe.Input.Complete(fault); } catch { }
+                try { pipe.Output.Complete(fault); } catch { }
 
                 if (fault != null && !_isShutdown)
                 {
@@ -339,14 +348,29 @@ private void Log(string message)
         }
 
         static Encoder s_sharedEncoder; // swapped in/out to avoid alloc on the public WriteResponse API
-        public static void WriteResponse(RedisClient client, PipeWriter output, RedisResult response)
+        public static ValueTask WriteResponseAsync(RedisClient client, PipeWriter output, RedisResult response)
         {
+            async ValueTask Awaited(ValueTask wwrite, Encoder eenc)
+            {
+                await wwrite;
+                Interlocked.Exchange(ref s_sharedEncoder, eenc);
+            }
             var enc = Interlocked.Exchange(ref s_sharedEncoder, null) ?? Encoding.UTF8.GetEncoder();
-            WriteResponse(client, output, response, enc);
+            var write = WriteResponseAsync(client, output, response, enc);
+            if (!write.IsCompletedSuccessfully) return Awaited(write, enc);
             Interlocked.Exchange(ref s_sharedEncoder, enc);
+            return default;
         }
-        internal static void WriteResponse(RedisClient client, PipeWriter output, RedisResult response, Encoder encoder)
+
+        internal static async ValueTask WriteResponseAsync(RedisClient client, PipeWriter output, RedisResult response, Encoder encoder)
+        {
+            void WritePrefix(PipeWriter ooutput, char pprefix)
             {
+                var span = ooutput.GetSpan(1);
+                span[0] = (byte)pprefix;
+                ooutput.Advance(1);
+            }
+
             if (response == null) return; // not actually a request (i.e. empty/whitespace request)
             if (client != null && client.ShouldSkipResponse()) return; // intentionally skipping the result
             char prefix;
@@ -361,9 +385,7 @@ internal static void WriteResponse(RedisClient client, PipeWriter output, RedisR
                 case ResultType.SimpleString:
                     prefix = '+';
                     BasicMessage:
-                    var span = output.GetSpan(1);
-                    span[0] = (byte)prefix;
-                    output.Advance(1);
+                    WritePrefix(output, prefix);
 
                     var val = response.AsString();
 
@@ -388,7 +410,9 @@ internal static void WriteResponse(RedisClient client, PipeWriter output, RedisR
                             var item = arr[i];
                             if (item == null)
                                 throw new InvalidOperationException("Array element cannot be null, index " + i);
-                            WriteResponse(null, output, item, encoder); // note: don't pass client down; this would impact SkipReplies
+
+                            // note: don't pass client down; this would impact SkipReplies
+                            await WriteResponseAsync(null, output, item, encoder);
                         }
                     }
                     break;
@@ -396,6 +420,7 @@ internal static void WriteResponse(RedisClient client, PipeWriter output, RedisR
                     throw new InvalidOperationException(
                         "Unexpected result type: " + response.Type);
             }
+            await output.FlushAsync();
         }
         public static bool TryParseRequest(ref ReadOnlySequence<byte> buffer, out RedisRequest request)
         {
@@ -411,32 +436,39 @@ public static bool TryParseRequest(ref ReadOnlySequence<byte> buffer, out RedisR
 
             return false;
         }
-        public bool TryProcessRequest(ref ReadOnlySequence<byte> buffer, RedisClient client, PipeWriter output)
+        public ValueTask<bool> TryProcessRequestAsync(ref ReadOnlySequence<byte> buffer, RedisClient client, PipeWriter output)
+        {
+            async ValueTask<bool> Awaited(ValueTask wwrite)
             {
+                await wwrite;
+                return true;
+            }
             if (!buffer.IsEmpty && TryParseRequest(ref buffer, out var request))
             {
                 RedisResult response;
                 try { response = Execute(client, request); }
                 finally { request.Recycle(); }
-                WriteResponse(client, output, response);
-                return true;
+                var write = WriteResponseAsync(client, output, response);
+                if (!write.IsCompletedSuccessfully) return Awaited(write);
+                return new ValueTask<bool>(true);
             }
-            return false;
+            return new ValueTask<bool>(false);
         }
 
-        private object ServerSyncLock => this;
+        protected object ServerSyncLock => this;
 
-        private long _commandsProcesed;
-        public long CommandsProcesed => _commandsProcesed;
+        private long _totalCommandsProcesed, _totalErrorCount;
+        public long TotalCommandsProcesed => _totalCommandsProcesed;
+        public long TotalErrorCount => _totalErrorCount;
 
         public RedisResult Execute(RedisClient client, RedisRequest request)
         {
             if (string.IsNullOrWhiteSpace(request.Command)) return null; // not a request
-            Interlocked.Increment(ref _commandsProcesed);
+            Interlocked.Increment(ref _totalCommandsProcesed);
             try
             {
                 RedisResult result;
-                if(_commands.TryGetValue(request.Command, out var cmd))
+                if (_commands.TryGetValue(request.Command, out var cmd))
                 {
                     request = request.AsCommand(cmd.Command); // fixup casing
                     if (cmd.HasSubCommands)
@@ -444,13 +476,13 @@ public RedisResult Execute(RedisClient client, RedisRequest request)
                         cmd = cmd.Resolve(request);
                         if (cmd.IsUnknown) return request.UnknownSubcommandOrArgumentCount();
                     }
-                    if(cmd.LockFree)
+                    if (cmd.LockFree)
                     {
                         result = cmd.Execute(client, request);
                     }
                     else
                     {
-                        lock(ServerSyncLock)
+                        lock (ServerSyncLock)
                         {
                             result = cmd.Execute(client, request);
                         }
@@ -462,6 +494,7 @@ public RedisResult Execute(RedisClient client, RedisRequest request)
                 }
 
                 if (result == null) Log($"missing command: '{request.Command}'");
+                else if (result.Type == ResultType.Error) Interlocked.Increment(ref _totalErrorCount);
                 return result ?? CommandNotFound(request.Command);
             }
             catch (NotSupportedException)
@@ -480,7 +513,7 @@ public RedisResult Execute(RedisClient client, RedisRequest request)
             }
             catch (Exception ex)
             {
-                if(!_isShutdown) Log(ex.Message);
+                if (!_isShutdown) Log(ex.Message);
                 return RedisResult.Create("ERR " + ex.Message, ResultType.Error);
             }
         }