Fix large string handling in CLUSTER|NODES and CLUSTER|SHARDS (#1858)

* fix CLUSTER|NODES and CLUSTER|SHARDS to handle larger string responses; another example of the issue we should cleanup up re. #1781

* fixes

* more fixes

* even more fixes; tricky to actually test this

Author: kevin-montrose

libs/cluster/Session/RespClusterBasicCommands.cs

@@ -2,6 +2,7 @@
 // Licensed under the MIT license.
 
 using System;
+using System.Buffers;
 using System.Diagnostics;
 using System.Text;
 using Garnet.common;
@@ -274,8 +275,8 @@ private bool NetworkClusterNodes(out bool invalidParameters)
             }
 
             var nodes = clusterProvider.clusterManager.CurrentConfig.GetClusterInfo(clusterProvider);
-            while (!RespWriteUtils.TryWriteAsciiBulkString(nodes, ref dcurr, dend))
-                SendAndReset();
+
+            WriteAsciiLargeRespString(nodes);
 
             return true;
         }
@@ -343,8 +344,8 @@ private bool NetworkClusterShards(out bool invalidParameters)
 
             var preferredType = clusterProvider.serverOptions.ClusterPreferredEndpointType;
             var shardsInfo = clusterProvider.clusterManager.CurrentConfig.GetShardsInfo(clusterProvider.clusterManager.clusterConnectionStore, preferredType);
-            while (!RespWriteUtils.TryWriteAsciiDirect(shardsInfo, ref dcurr, dend))
-                SendAndReset();
+
+            WriteLargeAsciiDirectString(shardsInfo);
 
             return true;
         }
@@ -523,5 +524,81 @@ private bool NetworkClusterPublish(out bool invalidParameters)
             clusterProvider.storeWrapper.subscribeBroker.Publish(parseState.GetArgSliceByRef(0), parseState.GetArgSliceByRef(1));
             return true;
         }
+
+
+        /// <summary>
+        /// Handle a potentially large already encoded RESP response, breaking it into pieces if necessary.
+        /// </summary>
+        private void WriteLargeAsciiDirectString(ReadOnlySpan<char> message)
+        {
+            // Attempt to write w/o any buffering
+            if (RespWriteUtils.TryWriteAsciiDirect(message, ref dcurr, dend))
+            {
+                return;
+            }
+
+            var buffer = ArrayPool<byte>.Shared.Rent(message.Length);
+            try
+            {
+                var len = Encoding.ASCII.GetBytes(message, buffer);
+                var remaining = buffer.AsSpan()[..len];
+
+                while (!remaining.IsEmpty)
+                {
+                    var space = Math.Min((int)(dend - dcurr), remaining.Length);
+                    _ = RespWriteUtils.TryWriteDirect(remaining[..space], ref dcurr, dend);
+
+                    SendAndReset();
+
+                    remaining = remaining[space..];
+                }
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
+
+        /// <summary>
+        /// Handle a potentially large bulk string by breaking it into pieces if necessary.
+        /// </summary>
+        private void WriteAsciiLargeRespString(ReadOnlySpan<char> message)
+        {
+            while (!RespWriteUtils.TryWriteBulkStringLength(message.Length, ref dcurr, dend))
+                SendAndReset();
+
+            // Attempt to write w/o any buffering
+            if (RespWriteUtils.TryWriteAsciiDirect(message, ref dcurr, dend))
+            {
+                while (!RespWriteUtils.TryWriteNewLine(ref dcurr, dend))
+                    SendAndReset();
+
+                return;
+            }
+
+            var buffer = ArrayPool<byte>.Shared.Rent(message.Length);
+            try
+            {
+                var len = Encoding.ASCII.GetBytes(message, buffer);
+                var remaining = buffer.AsSpan()[..len];
+
+                while (!remaining.IsEmpty)
+                {
+                    var space = Math.Min((int)(dend - dcurr), remaining.Length);
+                    _ = RespWriteUtils.TryWriteDirect(remaining[..space], ref dcurr, dend);
+
+                    SendAndReset();
+
+                    remaining = remaining[space..];
+                }
+
+                while (!RespWriteUtils.TryWriteNewLine(ref dcurr, dend))
+                    SendAndReset();
+            }
+            finally
+            {
+                ArrayPool<byte>.Shared.Return(buffer);
+            }
+        }
     }
 }

libs/common/RespWriteUtils.cs

@@ -321,14 +321,20 @@ public static bool TryWriteDirect<T>(ref T item, ref byte* curr, byte* end) wher
         /// Write length header of bulk string
         /// </summary>
         public static bool TryWriteBulkStringLength(ReadOnlySpan<byte> item, ref byte* curr, byte* end)
+        => TryWriteBulkStringLength(item.Length, ref curr, end);
+
+        /// <summary>
+        /// Write length header of bulk string
+        /// </summary>
+        public static bool TryWriteBulkStringLength(int len, ref byte* curr, byte* end)
         {
-            var itemDigits = NumUtils.CountDigits(item.Length);
+            var itemDigits = NumUtils.CountDigits(len);
             var totalLen = 1 + itemDigits + 2;
             if (totalLen > (int)(end - curr))
                 return false;
 
             *curr++ = (byte)'$';
-            NumUtils.WriteInt32(item.Length, itemDigits, ref curr);
+            NumUtils.WriteInt32(len, itemDigits, ref curr);
             WriteNewline(ref curr);
             return true;
         }