Skip to content

Pass RESP3 to modules and custom commands #1208

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
TalZaccai merged 7 commits into from
May 23, 2025
Merged

Pass RESP3 to modules and custom commands #1208

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
099b081
Let custom functions use RespMemoryWriter instead of IMemoryOwner<byt…
prvyk May 3, 2025
b49bf2d
Rename for clarity.
prvyk May 20, 2025
7a03832
Rename again.
prvyk May 20, 2025
3dfc08f
Clearer allocation minimum.
prvyk May 21, 2025
48bbc7e
Allow (re)allocating up to Array.MaxLength items so API will be as ca…
prvyk May 21, 2025
35dabd1
Merge branch 'main' into resp3modules
TalZaccai May 23, 2025
59a4fac
Merge branch 'main' into resp3modules
TalZaccai May 23, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
36 changes: 31 additions & 5 deletions libs/common/RespMemoryWriter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,7 +3,9 @@

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using System.Numerics;
using System.Runtime.CompilerServices;
using Tsavorite.core;
Expand Down Expand Up @@ -94,6 +96,21 @@ public void WriteBulkString(scoped ReadOnlySpan<byte> item)
ReallocateOutput(item.Length);
}

/// <summary>
/// Write bulk strings to memory.
/// </summary>
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void WriteBulkStrings(IEnumerable<byte[]> bulkStrings)
{
var stringLen = bulkStrings.Sum(x => x.Length);

while (!RespWriteUtils.TryWriteBulkString(bulkStrings, stringLen, ref curr, end))
{
var len = RespWriteUtils.GetBulkStringLength(stringLen);
ReallocateOutput(len);
}
}

/// <summary>
/// Writes the contents of <paramref name="span"/> as byte array to memory.
/// </summary>
Expand Down Expand Up @@ -285,6 +302,12 @@ public void WriteNewLine()
[MethodImpl(MethodImplOptions.AggressiveInlining)]
public void WriteNull()
{
// This is usually emitted by itself, so we can optimize allocation a bit.
if (end == curr)
{
Realloc(8);
}

if (resp3)
{
while (!RespWriteUtils.TryWriteResp3Null(ref curr, end))
Expand Down Expand Up @@ -379,30 +402,33 @@ public void WriteUtf8BulkString(ReadOnlySpan<char> chars)
/// Make sure at least totalLen bytes are allocated.
/// </summary>
/// <param name="totalLenHint"></param>
public void Realloc(int totalLenHint = 0)
public void Realloc(int totalLenHint)
{
var len = (int)(end - ptr);
if (totalLenHint <= len)
return;

ReallocateOutput(totalLenHint - len);
ReallocateOutput(totalLenHint - len, true);
}

[MethodImpl(MethodImplOptions.AggressiveInlining)]
private void ReallocateOutput(int extraLenHint = 0)
private void ReallocateOutput(int extraLenHint = 0, bool lowerMinimum = false)
{
var length = output.Length;

if (length < 1024)
{
length = 512;
if (!lowerMinimum)
length = 512;
else // Internally pool has a minimum of 16, no point allocating less.
length = 8;
}

checked
{
if (length < extraLenHint)
{
length = (int)BitOperations.RoundUpToPowerOf2((uint)extraLenHint + (uint)length); ;
length = (int)BitOperations.RoundUpToPowerOf2((uint)extraLenHint + (uint)length);
}
else
{
Expand Down
132 changes: 0 additions & 132 deletions libs/server/Custom/CustomCommandUtils.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,21 +2,11 @@
// Licensed under the MIT license.

using System;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Linq;
using Garnet.common;

namespace Garnet.server
{
internal static class CustomCommandUtils
{
/// <summary>
/// Shared memory pool used by functions
/// </summary>
private static MemoryPool<byte> MemoryPool => MemoryPool<byte>.Shared;

/// <summary>
/// Get first arg from input
/// </summary>
Expand Down Expand Up @@ -68,127 +58,5 @@ public static ReadOnlySpan<byte> GetNextArg(ref RawStringInput input, scoped ref
idx++;
return arg;
}

/// <summary>
/// Create output as bulk string, from given Span
/// </summary>
public static unsafe void WriteBulkString(ref (IMemoryOwner<byte>, int) output, Span<byte> bulkString)
{
// Get space for bulk string
var len = RespWriteUtils.GetBulkStringLength(bulkString.Length);
output.Item1 = MemoryPool.Rent(len);
output.Item2 = len;
fixed (byte* ptr = output.Item1.Memory.Span)
{
var curr = ptr;
// NOTE: Expected to always have enough space to write into pre-allocated buffer
var success = RespWriteUtils.TryWriteBulkString(bulkString, ref curr, ptr + len);
Debug.Assert(success, "Insufficient space in pre-allocated buffer");
}
}

/// <summary>
/// Create output as bulk string, from given Span
/// </summary>
public static unsafe void WriteBulkString(ref (IMemoryOwner<byte>, int) output, IEnumerable<byte[]> bulkStrings)
{
// Get space for bulk string
var stringLen = bulkStrings.Sum(x => x.Length);
var len = RespWriteUtils.GetBulkStringLength(stringLen);
output.Item1 = MemoryPool.Rent(len);
output.Item2 = len;
fixed (byte* ptr = output.Item1.Memory.Span)
{
var curr = ptr;
// NOTE: Expected to always have enough space to write into pre-allocated buffer
var success = RespWriteUtils.TryWriteBulkString(bulkStrings, stringLen, ref curr, ptr + len);
Debug.Assert(success, "Insufficient space in pre-allocated buffer");
}
}

/// <summary>
/// Create output as error message, from given string
/// </summary>
public static unsafe void WriteError(ref (IMemoryOwner<byte>, int) output, string errorMessage)
{
var bytes = System.Text.Encoding.ASCII.GetBytes(errorMessage);
WriteError(ref output, bytes);
}

/// <summary>
/// Create output as error message, from given byte span
/// </summary>
/// <param name="output">A tuple containing the memory owner and the length of written data</param>
/// <param name="errorMessage">The error message to write</param>
public static unsafe void WriteError(ref (IMemoryOwner<byte>, int) output, ReadOnlySpan<byte> errorMessage)
{
var bytes = errorMessage;
// Get space for error
var len = 1 + bytes.Length + 2;
output.Item1 = MemoryPool.Rent(len);
fixed (byte* ptr = output.Item1.Memory.Span)
{
var curr = ptr;
// NOTE: Expected to always have enough space to write into pre-allocated buffer
var success = RespWriteUtils.TryWriteError(bytes, ref curr, ptr + len);
Debug.Assert(success, "Insufficient space in pre-allocated buffer");
}
output.Item2 = len;
}

/// <summary>
/// Create null output as bulk string
/// </summary>
public static unsafe void WriteNullBulkString(ref (IMemoryOwner<byte>, int) output)
{
// Get space for null bulk string "$-1\r\n"
var len = 5;
output.Item1 = MemoryPool.Rent(len);
output.Item2 = len;
fixed (byte* ptr = output.Item1.Memory.Span)
{
var curr = ptr;
// NOTE: Expected to always have enough space to write into pre-allocated buffer
var success = RespWriteUtils.TryWriteNull(ref curr, ptr + len);
Debug.Assert(success, "Insufficient space in pre-allocated buffer");
}
}

/// <summary>
/// Create output as simple string, from given string
/// </summary>
public static unsafe void WriteSimpleString(ref (IMemoryOwner<byte>, int) output, string simpleString)
{
var bytes = System.Text.Encoding.ASCII.GetBytes(simpleString);
// Get space for simple string
var len = 1 + bytes.Length + 2;
output.Item1 = MemoryPool.Rent(len);
fixed (byte* ptr = output.Item1.Memory.Span)
{
var curr = ptr;
// NOTE: Expected to always have enough space to write into pre-allocated buffer
var success = RespWriteUtils.TryWriteSimpleString(bytes, ref curr, ptr + len);
Debug.Assert(success, "Insufficient space in pre-allocated buffer");
}
output.Item2 = len;
}

/// <summary>
/// Writes bytes directly into a rented memory buffer without any encoding transformation.
/// </summary>
/// <param name="output">A tuple containing the memory owner and the length of written data.</param>
/// <param name="bytes">The source bytes to write.</param>
public static unsafe void WriteDirect(ref (IMemoryOwner<byte>, int) output, ReadOnlySpan<byte> bytes)
{
output.Item1 = MemoryPool.Rent(bytes.Length);
fixed (byte* ptr = output.Item1.Memory.Span)
{
var curr = ptr;
// NOTE: Expected to always have enough space to write into pre-allocated buffer
var success = RespWriteUtils.TryWriteDirect(bytes, ref curr, ptr + bytes.Length);
Debug.Assert(success, "Insufficient space in pre-allocated buffer");
}
output.Item2 = bytes.Length;
}
}
}
Loading