diff --git a/src/StackExchange.Redis/Message.cs b/src/StackExchange.Redis/Message.cs index 837651106..c37c56751 100644 --- a/src/StackExchange.Redis/Message.cs +++ b/src/StackExchange.Redis/Message.cs @@ -536,7 +536,12 @@ bool ICompletable.TryComplete(bool isAsync) return true; } - public void Complete(PhysicalConnection? connection) + // virtual so a subclass owning a rented request buffer (e.g. RenderedArgs) can recycle it here: + // every path that terminally finishes a message - normal completion, high-integrity validation, + // and the RecordConnectionFailed drain - runs through this (via SetExceptionAndComplete, for the + // last one), and PrepareToResend/caller-side timeouts never call it at all, which is exactly the + // "message is definitely never touched again" boundary the buffer needs. + public virtual void Complete(PhysicalConnection? connection) { // Ensure we can never call Complete on the same resultBox from two threads by grabbing it now var currBox = Interlocked.Exchange(ref resultBox, null); diff --git a/src/StackExchange.Redis/RedisDatabase.cs b/src/StackExchange.Redis/RedisDatabase.cs index a5148ee57..0f3ed8c6f 100644 --- a/src/StackExchange.Redis/RedisDatabase.cs +++ b/src/StackExchange.Redis/RedisDatabase.cs @@ -2007,7 +2007,7 @@ public Task PublishAsync(RedisChannel channel, RedisValue message, Command public RespResult ExecuteResp(string command, ReadOnlyMemory args, CommandFlags flags = CommandFlags.None) { - var msg = new ExecMessage(multiplexer?.CommandMap, Database, flags, command, args); + var msg = new ExecMessage(multiplexer?.CommandMap, Database, flags, command, args, multiplexer?.RawConfig?.RequestBufferPool); return ExecuteSync(msg, ResultProcessor.RespResult)!; } @@ -2022,7 +2022,7 @@ public RedisResult Execute(string command, ICollection args, CommandFlag public Task ExecuteRespAsync(string command, ReadOnlyMemory args, CommandFlags flags = CommandFlags.None) { - var msg = new ExecMessage(multiplexer?.CommandMap, Database, flags, command, args); + var msg = new ExecMessage(multiplexer?.CommandMap, Database, flags, command, args, multiplexer?.RawConfig?.RequestBufferPool); return ExecuteAsync(msg, ResultProcessor.RespResult, defaultValue: RespResult.NullReply); } @@ -2038,7 +2038,7 @@ public Task ExecuteAsync(string command, ICollection? args, public RespResult ScriptEvaluateResp(string script, ReadOnlyMemory keys, ReadOnlyMemory values, CommandFlags flags = CommandFlags.None) { var command = ResultProcessor.ScriptLoadProcessor.IsSHA1(script) ? RedisCommand.EVALSHA : RedisCommand.EVAL; - var msg = new ScriptEvalMessage(Database, flags, command, script, keys, values); + var msg = new ScriptEvalMessage(Database, flags, command, script, keys, values, multiplexer?.RawConfig?.RequestBufferPool); try { return ExecuteSync(msg, ResultProcessor.RespResult)!; @@ -2084,7 +2084,7 @@ public RedisResult ScriptEvaluate(LoadedLuaScript script, object? parameters = n public async Task ScriptEvaluateRespAsync(string script, ReadOnlyMemory keys, ReadOnlyMemory values, CommandFlags flags = CommandFlags.None) { var command = ResultProcessor.ScriptLoadProcessor.IsSHA1(script) ? RedisCommand.EVALSHA : RedisCommand.EVAL; - var msg = new ScriptEvalMessage(Database, flags, command, script, keys, values); + var msg = new ScriptEvalMessage(Database, flags, command, script, keys, values, multiplexer?.RawConfig?.RequestBufferPool); try { @@ -2158,7 +2158,7 @@ public RespResult ScriptEvaluateReadOnlyResp(string script, ReadOnlyMemory ScriptEvaluateReadOnlyRespAsync(string script, Rea multiplexer.CommandMap, ResultProcessor.ScriptLoadProcessor.IsSHA1(script) ? RedisCommand.EVALSHA_RO : RedisCommand.EVAL_RO, ref flags); - var msg = new ScriptEvalMessage(Database, flags, command, script, keys, values); + var msg = new ScriptEvalMessage(Database, flags, command, script, keys, values, multiplexer?.RawConfig?.RequestBufferPool); try { return await ExecuteAsync(msg, ResultProcessor.RespResult, defaultValue: RespResult.NullReply).ForAwait(); @@ -6045,7 +6045,11 @@ protected override bool SetResultCore(PhysicalConnection connection, Message mes internal sealed class ExecMessage : Message { - private readonly ReadOnlyMemory _args; + // not readonly: RenderedArgs.Recycle swaps the buffer out, and cannot do that through a + // defensive copy - see RenderedArgs + private RenderedArgs _args; + private readonly bool _hasSubCommand; + private readonly SubCommand _subCommand; private string _unknownCommand; private static int RemoveDbIfNotRequired(int suggestedDb, string adhocCommand, out RedisCommand knownCommand) @@ -6065,7 +6069,7 @@ private static int RemoveDbIfNotRequired(int suggestedDb, string adhocCommand, o return suggestedDb; } - public ExecMessage(CommandMap? map, int db, CommandFlags flags, string command, ReadOnlyMemory args) + public ExecMessage(CommandMap? map, int db, CommandFlags flags, string command, ReadOnlyMemory args, MemoryPool? pool) : base(RemoveDbIfNotRequired(db, command, out var knownCommand), flags, knownCommand) { if (args.Length >= MessageWriter.REDIS_MAX_ARGS) // using >= here because we will be adding 1 for the command itself (which is an arg for the purposes of the multi-bulk protocol) @@ -6088,71 +6092,59 @@ public ExecMessage(CommandMap? map, int db, CommandFlags flags, string command, { throw ExceptionFactory.CommandDisabled(command); } - _args = args; + + // resolved now, while we still have the arguments as values; only the first one is a + // sub-command candidate, which is what the old write-time loop amounted to + if (args.Length != 0) + { + ref readonly var first = ref args.Span[0]; + if (first.IsValue) _hasSubCommand = SubCommandMetadata.TryGetSubCommand(first.Value, out _subCommand); + } + + _args = RenderedArgs.Create(args.Span, pool); } protected override void WriteImpl(in MessageWriter writer) { if (Command is RedisCommand.UNKNOWN) { - writer.WriteHeader(_unknownCommand, _args.Length); + writer.WriteHeader(_unknownCommand, _args.Count); } else { - writer.WriteHeader(Command, _args.Length); - } - foreach (ref readonly var arg in _args.Span) - { - if (arg.IsKey) - { - writer.Write(arg.Key); - } - else if (arg.IsValue) - { - writer.WriteBulkString(arg.Value); - } - else - { - Debug.Assert(arg.IsNull); - throw new InvalidOperationException("A null is not valid in this context"); - } + writer.WriteHeader(Command, _args.Count); } + + // keys already carry their prefix, and a key and a value are the same thing on the wire + _args.WriteTo(writer); } public override string CommandString => Command is RedisCommand.UNKNOWN ? _unknownCommand : base.CommandString; public override string CommandAndKey => CommandString; public override int GetHashSlot(ServerSelectionStrategy serverSelectionStrategy) + => _args.GetHashSlot(serverSelectionStrategy); + + public override int ArgCount => _args.Count; + + protected override bool TryGetSubCommand(out SubCommand subCommand) { - int slot = ServerSelectionStrategy.NoSlot; - foreach (ref readonly var arg in _args.Span) - { - var key = arg.Key; - if (!key.IsNull) - { - slot = serverSelectionStrategy.CombineSlot(slot, key); - } - } - return slot; + // resolved in the constructor: by now the arguments are rendered bytes, and this needs + // them as values + subCommand = _hasSubCommand ? _subCommand : SubCommand.Unknown; + return _hasSubCommand; } - public override int ArgCount => _args.Length; - protected override bool TryGetSubCommand(out SubCommand subCommand) + public override void Complete(PhysicalConnection? connection) { - // the sub-command (if any) is the first argument after the command itself, - // e.g. CLIENT [GETNAME]; ad-hoc Execute args are boxed objects, so normalize - // the first one to a RedisValue before probing it against the known sub-commands - foreach (ref readonly var arg in _args.Span) - { - var value = arg.Value; - if (!value.IsNull) - { - return SubCommandMetadata.TryGetSubCommand(value, out subCommand); - } - break; // only the first argument is a sub-command candidate - } - subCommand = SubCommand.Unknown; - return false; + // Status flips to Sent strictly *after* WriteImpl returns (see WriteMessageToServerInsideWriteLock), + // but a message is enqueued into _writtenAwaitingResponse - making it eligible for the async-timeout + // heartbeat's SetExceptionAndComplete - *before* that write actually happens. Recycling on that + // premature completion would race an in-flight WriteImpl reading the same buffer. If the write did + // land, the real reply's own later Complete() call recycles once Status has caught up; if it never + // gets sent at all, the buffer leaks rather than risking corruption - the lesser evil. + if (Status == CommandStatus.Sent) RenderedArgs.Recycle(ref _args); + base.Complete(connection); } } @@ -6297,28 +6289,24 @@ private static bool IsReadOnlyScript(RedisCommand command) private sealed class ScriptEvalMessage : Message, IMultiMessage { - private readonly ReadOnlyMemory _keys; - private readonly ReadOnlyMemory _values; + // not readonly: RenderedArgs.Recycle swaps the buffer out, and cannot do that through a + // defensive copy - see RenderedArgs. The script itself stays out of the buffer: it is needed + // intact for hash lookup and SCRIPT LOAD, and being a string it carries no lifetime hazard. + private RenderedArgs _args; + private readonly int _keyCount; private readonly string _script; private byte[]? asciiHash; private bool useReadOnly; - public ScriptEvalMessage(int db, CommandFlags flags, RedisCommand command, string script, ReadOnlyMemory keys, ReadOnlyMemory values) + public ScriptEvalMessage(int db, CommandFlags flags, RedisCommand command, string script, ReadOnlyMemory keys, ReadOnlyMemory values, MemoryPool? pool) : base(db, flags, command) { _script = script ?? throw new ArgumentNullException(nameof(script)); - _keys = keys; - _values = values; + _keyCount = keys.Length; + _args = RenderedArgs.Create(keys.Span, values.Span, pool); } public override int GetHashSlot(ServerSelectionStrategy serverSelectionStrategy) - { - int slot = ServerSelectionStrategy.NoSlot; - foreach (ref readonly var key in _keys.Span) - { - slot = serverSelectionStrategy.CombineSlot(slot, key); - } - return slot; - } + => _args.GetHashSlot(serverSelectionStrategy); public IEnumerable GetMessages(PhysicalConnection connection) { @@ -6348,28 +6336,29 @@ protected override void WriteImpl(in MessageWriter writer) { if (asciiHash != null) { - writer.WriteHeader(useReadOnly ? RedisCommand.EVALSHA_RO : RedisCommand.EVALSHA, 2 + _keys.Length + _values.Length); + writer.WriteHeader(useReadOnly ? RedisCommand.EVALSHA_RO : RedisCommand.EVALSHA, ArgCount); writer.WriteBulkString(asciiHash); } else { - writer.WriteHeader(useReadOnly ? RedisCommand.EVAL_RO : RedisCommand.EVAL, 2 + _keys.Length + _values.Length); + writer.WriteHeader(useReadOnly ? RedisCommand.EVAL_RO : RedisCommand.EVAL, ArgCount); writer.WriteBulkString(_script); } - writer.WriteBulkString(_keys.Length); + writer.WriteBulkString(_keyCount); - foreach (ref readonly var key in _keys.Span) - { - writer.Write(key); - } + // rendered keys-then-values, in the order EVAL wants them; keys already carry their prefix + _args.WriteTo(writer); + } - foreach (ref readonly var value in _values.Span) - { - writer.WriteBulkString(value); - } + public override int ArgCount => 2 + _args.Count; + + public override void Complete(PhysicalConnection? connection) + { + // see ExecMessage.Complete for why this is gated on Status rather than unconditional + if (Status == CommandStatus.Sent) RenderedArgs.Recycle(ref _args); + base.Complete(connection); } - public override int ArgCount => 2 + _keys.Length + _values.Length; } private sealed class ScriptEvaluateMessage : Message, IMultiMessage diff --git a/src/StackExchange.Redis/RenderedArgs.cs b/src/StackExchange.Redis/RenderedArgs.cs new file mode 100644 index 000000000..df455da3d --- /dev/null +++ b/src/StackExchange.Redis/RenderedArgs.cs @@ -0,0 +1,228 @@ +using System; +using System.Buffers; +using System.Diagnostics; +using System.Runtime.CompilerServices; +using System.Runtime.InteropServices; +using System.Threading; + +namespace StackExchange.Redis +{ + /// + /// The keys and values of a request, rendered once into a single pooled buffer at the point of the + /// call, so that the request no longer refers to any memory the caller owns. + /// + /// + /// + /// Callers can pass keys and values backed by their own arrays - and are encouraged to, for the sake + /// of allocation - but the request is not necessarily written before the call returns: it may be + /// queued in the backlog, or in a batch, and fire-and-forget callers get no completion signal at all. + /// Rendering here means the caller's buffers are theirs again the moment the call returns, and it also + /// moves the formatting off the writer thread, where it would otherwise be done while holding the + /// single-writer lock. + /// + /// + /// Entries are laid out back to back as a 4-byte little-endian length followed by that many payload + /// bytes. A negative length marks a key, with ~length giving the real size; the complement + /// rather than negation so that a zero-length key is still distinguishable from a zero-length value. + /// + /// + /// This type owns a pooled buffer, so it must live in exactly one place and must never be copied. + /// Copying it duplicates the ownership and gets the buffer returned to the pool twice, which is a + /// silent corruption rather than a loud failure. Note also that the field holding it must not be + /// readonly - see . + /// + /// + internal struct RenderedArgs + { + // deliberately not readonly: Recycle swaps this out atomically. Either a byte[] from + // ArrayPool.Shared or an IMemoryOwner from the configured RequestBufferPool - the + // same shape RespResult and Lease use, so that a caller who supplies a pool gets it honoured + // on the request side too. + private object _buffer; + private int _count, _length; + + /// The number of entries rendered. + public readonly int Count => _count; + + private const int PrefixLength = sizeof(int); + + /// + /// Render the arguments of an ad-hoc command, which may be a mix of keys and values. + /// + public static RenderedArgs Create(ReadOnlySpan args, MemoryPool? pool) + { + int total = 0; + foreach (ref readonly var arg in args) + { + // rejected here rather than at write time: the caller finds out on the call that made + // the mistake, instead of from a background writer some time later + if (arg.IsNull) throw new InvalidOperationException("A null is not valid in this context"); + total += PrefixLength + (arg.IsKey ? arg.Key.TotalLength() : arg.Value.GetByteCount()); + } + + var result = Rent(args.Length, total, pool); + var target = result.Buffer; + foreach (ref readonly var arg in args) + { + target = arg.IsKey ? WriteKey(target, arg.Key) : WriteValue(target, arg.Value); + } + + Debug.Assert(target.IsEmpty, "should have filled the buffer exactly"); + return result; + } + + /// + /// Render the keys and values of a script invocation; keys are emitted first, as EVAL expects. + /// + public static RenderedArgs Create(ReadOnlySpan keys, ReadOnlySpan values, MemoryPool? pool) + { + int total = 0; + foreach (ref readonly var key in keys) + { + key.AssertNotNull(); + total += PrefixLength + key.TotalLength(); + } + foreach (ref readonly var value in values) + { + value.AssertNotNull(); + total += PrefixLength + value.GetByteCount(); + } + + var result = Rent(keys.Length + values.Length, total, pool); + var target = result.Buffer; + foreach (ref readonly var key in keys) target = WriteKey(target, key); + foreach (ref readonly var value in values) target = WriteValue(target, value); + + Debug.Assert(target.IsEmpty, "should have filled the buffer exactly"); + return result; + } + + private static RenderedArgs Rent(int count, int length, MemoryPool? pool) => new RenderedArgs + { + _buffer = length == 0 ? Array.Empty() + : pool is null ? ArrayPool.Shared.Rent(length) : pool.Rent(length), + _count = count, + _length = length, + }; + + // the rented buffer is usually larger than we asked for; only the used span is ours + private readonly Span Buffer => _buffer switch + { + byte[] arr => new Span(arr, 0, _length), + IMemoryOwner owner => owner.Memory.Span.Slice(0, _length), + _ => default, + }; + + private static Span WriteKey(Span target, scoped in RedisKey key) + { + var length = key.TotalLength(); + Unsafe.WriteUnaligned(ref target[0], ~length); + var written = key.CopyTo(target.Slice(PrefixLength)); + Debug.Assert(written == length, "key length disagreed with itself"); + return target.Slice(PrefixLength + length); + } + + private static Span WriteValue(Span target, scoped in RedisValue value) + { + var length = value.GetByteCount(); + Unsafe.WriteUnaligned(ref target[0], length); + var written = value.CopyTo(target.Slice(PrefixLength)); + Debug.Assert(written == length, "value length disagreed with itself"); + return target.Slice(PrefixLength + length); + } + + /// + /// Write every entry as a RESP bulk string, in the order they were rendered. + /// + /// + /// Keys and values are indistinguishable on the wire - the flag exists for slot routing, not for + /// framing - so this writes them identically. + /// + public readonly void WriteTo(in MessageWriter writer) + { + var iter = GetEnumerator(); + while (iter.MoveNext()) + { + writer.WriteBulkString(iter.Current); + } + } + + /// + /// The combined cluster slot of the keys, or when + /// there are none, or when they disagree. + /// + /// + /// Values are skipped: only arguments the caller declared as keys take part in routing. Note this + /// hashes the rendered bytes directly, where has + /// to copy the key out first - we already have exactly the bytes it would have produced. + /// + public readonly int GetHashSlot(ServerSelectionStrategy strategy) + { + if (strategy.ServerType is ServerType.Standalone) return ServerSelectionStrategy.NoSlot; + + var slot = ServerSelectionStrategy.NoSlot; + var iter = GetEnumerator(); + while (iter.MoveNext()) + { + if (!iter.IsKey) continue; + slot = ServerSelectionStrategy.CombineSlot(slot, ServerSelectionStrategy.GetClusterSlot(iter.Current)); + } + return slot; + } + + /// + /// Walks the rendered entries in order. + /// + public readonly Enumerator GetEnumerator() => new Enumerator(Buffer); + + /// + /// Return the buffer to the pool; safe to call repeatedly, and from multiple threads - only the + /// first caller sees the buffer. + /// + /// + /// Taken by ref on purpose. As an instance method this would compile perfectly happily + /// against a readonly field and silently operate on a defensive copy, leaving the real + /// buffer stranded; as a ref parameter, that same mistake is CS0192 at build time. + /// + public static void Recycle(ref RenderedArgs args) + { + var buffer = Interlocked.Exchange(ref args._buffer, Array.Empty()); + args._length = args._count = 0; + + // null when never rendered, empty when already recycled or nothing to render + if (buffer is byte[] { Length: > 0 } arr) ArrayPool.Shared.Return(arr); + else if (buffer is IMemoryOwner owner) owner.Dispose(); + } + + /// Walks rendered entries. + internal ref struct Enumerator(ReadOnlySpan remaining) + { + // read-only: walking never writes to the rendered buffer, it only slices through it + private ReadOnlySpan _remaining = remaining; + + /// The payload of the current entry. + public ReadOnlySpan Current { get; private set; } + + /// Whether the current entry was supplied as a key rather than a value. + public bool IsKey { get; private set; } + + /// Move to the next entry. + public bool MoveNext() + { + if (_remaining.IsEmpty) + { + Current = default; + IsKey = false; + return false; + } + + var prefix = Unsafe.ReadUnaligned(ref MemoryMarshal.GetReference(_remaining)); + IsKey = prefix < 0; + var length = IsKey ? ~prefix : prefix; + Current = _remaining.Slice(PrefixLength, length); + _remaining = _remaining.Slice(PrefixLength + length); + return true; + } + } + } +} diff --git a/tests/StackExchange.Redis.Tests/RenderedArgsLeakTests.cs b/tests/StackExchange.Redis.Tests/RenderedArgsLeakTests.cs new file mode 100644 index 000000000..500f6f0f2 --- /dev/null +++ b/tests/StackExchange.Redis.Tests/RenderedArgsLeakTests.cs @@ -0,0 +1,98 @@ +using System; +using System.Buffers; +using System.Threading; +using System.Threading.Tasks; +using Xunit; + +namespace StackExchange.Redis.Tests; + +[RunPerProtocol] +public class RenderedArgsLeakTests(ITestOutputHelper output) : TestBase(output) +{ + // end-to-end companion to the codec-level RenderedArgsTests: drives real ExecuteResp/ScriptEvaluateResp + // traffic against a live server and confirms the rented request buffers actually come back - normal + // completion, not just the isolated Recycle() calls the unit tests exercise directly. + [Fact] + public async Task ExecuteRespAndScriptEvaluateResp_DoNotLeakRentedBuffers() + { + const int Iterations = 500; + var pool = new CountingPool(); + var options = new ConfigurationOptions + { + EndPoints = { TestConfig.Current.PrimaryServerAndPort }, + RequestBufferPool = pool, + Protocol = TestContext.Current.GetProtocol(), + }; + + RedisKey key = Me(); + await using (var conn = await ConnectionMultiplexer.ConnectAsync(options)) + { + var db = conn.GetDatabase(); + for (int i = 0; i < Iterations; i++) + { + using var setResult = await db.ExecuteRespAsync("SET", new RedisKeyOrValue[] { key, (RedisValue)i }); + using var getResult = await db.ScriptEvaluateRespAsync("return redis.call('get', KEYS[1])", new RedisKey[] { key }, default); + Assert.Equal(i.ToString(), (string?)getResult.ReadScalar().ReadRedisValue()); + } + } + + // after full teardown, every rented buffer should have come back - modulo a small, stable + // residual for whatever the connection's own write/read buffering happened to be holding at the + // moment of measurement, not a per-call leak (see PR #3211 discussion: ballpark 500 calls -> ~550 + // rents [500 for the request buffers, the rest for IO], with almost all of them returned). + var outstanding = pool.Rented - pool.Returned; + Output.WriteLine($"rented={pool.Rented}, returned={pool.Returned}, outstanding={outstanding}"); + Assert.True(pool.Rented >= Iterations * 2, $"expected at least one rent per call, got {pool.Rented} for {Iterations} iterations"); + Assert.True(outstanding <= 5, $"expected only a small residual, got {outstanding} outstanding (rented={pool.Rented}, returned={pool.Returned})"); + } + + // DIAGNOSTIC: same loop, but through the *old*, pre-RenderedArgs ScriptEvaluateAsync path, to determine + // whether the "keys > args" corruption seen under full-suite load is specific to the new codec or a + // pre-existing race in the write path that predates this PR entirely. + [Fact] + public async Task OldScriptEvaluateAsync_DoesNotCorruptUnderLoad() + { + const int Iterations = 500; + RedisKey key = Me(); + await using var conn = await ConnectionMultiplexer.ConnectAsync(new ConfigurationOptions + { + EndPoints = { TestConfig.Current.PrimaryServerAndPort }, + Protocol = TestContext.Current.GetProtocol(), + }); + var db = conn.GetDatabase(); + for (int i = 0; i < Iterations; i++) + { + await db.StringSetAsync(key, i); + var result = await db.ScriptEvaluateAsync("return redis.call('get', KEYS[1])", new RedisKey[] { key }, default); + Assert.Equal(i.ToString(), (string?)result); + } + } + + private sealed class CountingPool : MemoryPool + { + private int _rented, _returned; + public int Rented => _rented; + public int Returned => _returned; + public override int MaxBufferSize => Shared.MaxBufferSize; + + public override IMemoryOwner Rent(int minBufferSize = -1) + { + Interlocked.Increment(ref _rented); + return new Owner(this, Shared.Rent(minBufferSize)); + } + + protected override void Dispose(bool disposing) + { + } + + private sealed class Owner(CountingPool pool, IMemoryOwner inner) : IMemoryOwner + { + public Memory Memory => inner.Memory; + public void Dispose() + { + Interlocked.Increment(ref pool._returned); + inner.Dispose(); + } + } + } +} diff --git a/tests/StackExchange.Redis.Tests/RenderedArgsTests.cs b/tests/StackExchange.Redis.Tests/RenderedArgsTests.cs new file mode 100644 index 000000000..ecbf042e7 --- /dev/null +++ b/tests/StackExchange.Redis.Tests/RenderedArgsTests.cs @@ -0,0 +1,280 @@ +using System; +using System.Buffers; +using System.Collections.Generic; +using System.Linq; +using System.Text; +using Xunit; + +namespace StackExchange.Redis.Tests; + +/// +/// The codec only: entries round-trip, keys and values stay distinguishable, and the buffer goes back to +/// the pool exactly once however many times it is asked to. +/// +public class RenderedArgsTests +{ + private static List<(bool IsKey, string Payload)> Drain(in RenderedArgs args) + { + var result = new List<(bool, string)>(); + var iter = args.GetEnumerator(); + while (iter.MoveNext()) + { + result.Add((iter.IsKey, Encoding.UTF8.GetString(iter.Current))); + } + return result; + } + + [Fact] + public void MixedKeysAndValuesRoundTrip() + { + RedisKeyOrValue[] input = [(RedisKey)"k1", (RedisValue)"v1", (RedisKey)"k2", (RedisValue)123]; + var args = RenderedArgs.Create(input, pool: null); + try + { + Assert.Equal(4, args.Count); + Assert.Equal( + [(true, "k1"), (false, "v1"), (true, "k2"), (false, "123")], + Drain(in args)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + [Fact] + public void ScriptFormPutsKeysFirst() + { + RedisKey[] keys = ["a", "bb"]; + RedisValue[] values = ["x", "yy", "zzz"]; + var args = RenderedArgs.Create(keys, values, pool: null); + try + { + Assert.Equal(5, args.Count); + Assert.Equal( + [(true, "a"), (true, "bb"), (false, "x"), (false, "yy"), (false, "zzz")], + Drain(in args)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + [Fact] + public void ZeroLengthKeyAndValueStayDistinguishable() + { + // this is why the prefix uses ~length rather than -length: negation cannot tell a zero-length + // key from a zero-length value, because -0 == 0 + RedisKeyOrValue[] input = [(RedisKey)"", (RedisValue)""]; + var args = RenderedArgs.Create(input, pool: null); + try + { + Assert.Equal([(true, ""), (false, "")], Drain(in args)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + [Fact] + public void EmptyInputRendersNothingAndRecyclesCleanly() + { + var args = RenderedArgs.Create([], pool: null); + Assert.Equal(0, args.Count); + Assert.Empty(Drain(in args)); + RenderedArgs.Recycle(ref args); + RenderedArgs.Recycle(ref args); + } + + [Fact] + public void DefaultInstanceIsSafeToRecycle() + { + // a message that fails before its arguments are ever rendered still has to be tidied up + RenderedArgs args = default; + Assert.Equal(0, args.Count); + RenderedArgs.Recycle(ref args); + } + + [Fact] + public void RecycleIsOnceOnly() + { + var pool = new CountingPool(); + var args = RenderedArgs.Create([(RedisValue)"payload"], pool); + Assert.Equal(1, pool.Rented); + + RenderedArgs.Recycle(ref args); + RenderedArgs.Recycle(ref args); + RenderedArgs.Recycle(ref args); + + Assert.Equal(1, pool.Returned); // and not three + Assert.Equal(0, args.Count); + Assert.Empty(Drain(in args)); // reading after recycling is empty rather than a fault + } + + [Fact] + public void ConfiguredPoolIsUsedInsteadOfTheSharedArrayPool() + { + var pool = new CountingPool(); + var args = RenderedArgs.Create([(RedisKey)"key", (RedisValue)"value"], pool); + try + { + Assert.Equal(1, pool.Rented); + Assert.Equal([(true, "key"), (false, "value")], Drain(in args)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + Assert.Equal(1, pool.Returned); + } + + [Fact] + public void LargePayloadsSurviveTheLengthPrefix() + { + var big = new string('x', 100_000); + var args = RenderedArgs.Create([(RedisValue)big], pool: null); + try + { + var drained = Drain(in args); + Assert.Single(drained); + Assert.Equal(big, drained[0].Payload); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + private static byte[] Write(in RenderedArgs args) + { + var writer = new MessageWriter(null, CommandMap.Default, MessageWriter.BlockBuffer); + ReadOnlyMemory payload = default; + try + { + args.WriteTo(writer); + payload = MessageWriter.FlushBlockBuffer(); + return payload.Span.ToArray(); + } + catch + { + MessageWriter.RevertBlockBuffer(); + throw; + } + finally + { + MessageWriter.ReleaseBlockBuffer(payload); + } + } + + [Fact] + public void WritesEachEntryAsABulkString() + { + var args = RenderedArgs.Create([(RedisKey)"key", (RedisValue)"value", (RedisValue)42], pool: null); + try + { + // keys and values are indistinguishable on the wire; the flag is only for routing + Assert.Equal("$3\r\nkey\r\n$5\r\nvalue\r\n$2\r\n42\r\n", Encoding.UTF8.GetString(Write(in args))); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + [Fact] + public void WritesNothingForNoArguments() + { + var args = RenderedArgs.Create([], pool: null); + try + { + Assert.Empty(Write(in args)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + [Theory] + [InlineData("abc")] // single key + [InlineData("abc,abc")] // same key twice - one slot + [InlineData("{tag}a,{tag}b")] // hash tags - still one slot + [InlineData("one,two")] // disagreeing keys - MultipleSlots + [InlineData("")] // no keys at all + public void HashSlotAgreesWithThePerKeyPath(string keyNames) + { + var strategy = new ServerSelectionStrategy(null) { ServerType = ServerType.Cluster }; + var keys = keyNames.Length == 0 + ? [] + : keyNames.Split(',').Select(x => (RedisKey)x).ToArray(); + + // what the existing per-key routing would decide + var expected = ServerSelectionStrategy.NoSlot; + foreach (var key in keys) expected = strategy.CombineSlot(expected, key); + + var args = RenderedArgs.Create(keys, [(RedisValue)"ignored", (RedisValue)"also ignored"], pool: null); + try + { + Assert.Equal(expected, args.GetHashSlot(strategy)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + [Fact] + public void HashSlotIgnoresValuesEvenWhenTheyLookLikeKeys() + { + var strategy = new ServerSelectionStrategy(null) { ServerType = ServerType.Cluster }; + + var keyed = RenderedArgs.Create([(RedisKey)"abc"], pool: null); + var valued = RenderedArgs.Create([(RedisValue)"abc"], pool: null); + try + { + Assert.Equal(strategy.HashSlot((RedisKey)"abc"), keyed.GetHashSlot(strategy)); + Assert.Equal(ServerSelectionStrategy.NoSlot, valued.GetHashSlot(strategy)); + } + finally + { + RenderedArgs.Recycle(ref keyed); + RenderedArgs.Recycle(ref valued); + } + } + + [Fact] + public void HashSlotIsNoSlotForStandalone() + { + var strategy = new ServerSelectionStrategy(null) { ServerType = ServerType.Standalone }; + var args = RenderedArgs.Create([(RedisKey)"abc"], pool: null); + try + { + Assert.Equal(ServerSelectionStrategy.NoSlot, args.GetHashSlot(strategy)); + } + finally + { + RenderedArgs.Recycle(ref args); + } + } + + private sealed class CountingPool : MemoryPool + { + public int Rented { get; private set; } + public int Returned { get; private set; } + public override int MaxBufferSize => Shared.MaxBufferSize; + public override IMemoryOwner Rent(int minBufferSize = -1) + { + Rented++; + return new Owner(this, Shared.Rent(minBufferSize)); + } + protected override void Dispose(bool disposing) { } + + private sealed class Owner(CountingPool pool, IMemoryOwner inner) : IMemoryOwner + { + public Memory Memory => inner.Memory; + public void Dispose() { pool.Returned++; inner.Dispose(); } + } + } +} \ No newline at end of file