Skip to content

Reduce allocation in random set index picking #1242

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

Open
wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Reduce allocation in random set index picking #1242

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
366230e
Reduce allocation in random set index picking
PaulusParssinen Jun 6, 2025
b30f4a5
Remove method level duplication of the stackalloc threshold constant
PaulusParssinen Jun 6, 2025
917fc16
Address PR feedback
PaulusParssinen Jun 11, 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
53 changes: 25 additions & 28 deletions libs/common/RandomUtils.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,20 +18,24 @@ public static class RandomUtils
/// Pick k indexes from a collection of n items
/// </summary>
/// <param name="n">Number of items in the collection</param>
/// <param name="k">Number of items to pick</param>
/// <param name="indexes">Span of indexes to pick.</param>
/// <param name="seed">Random seed</param>
/// <param name="distinct">Whether items returned should be distinct (default: true)</param>
/// <returns>K indexes picked</returns>
public static Span<int> PickKRandomIndexes(int n, int k, int seed, bool distinct = true)
public static void PickKRandomIndexes(int n, Span<int> indexes, int seed, bool distinct = true)
{
if (k < 0) throw new ArgumentOutOfRangeException(nameof(k));
if (n < 0) throw new ArgumentOutOfRangeException(nameof(n));
if (distinct && k > n) throw new ArgumentException(
$"{nameof(k)} cannot be larger than {nameof(n)} when indexes should be distinct.");
ArgumentOutOfRangeException.ThrowIfNegative(n);
if (distinct && indexes.Length > n) throw new ArgumentException(
$"{nameof(indexes)} cannot be larger than {nameof(n)} when indexes should be distinct.");

return !distinct || (double)k / n < KOverNThreshold
? PickKRandomIndexesIteratively(n, k, seed, distinct)
: PickKRandomDistinctIndexesWithShuffle(n, k, seed);
if (!distinct || (double)indexes.Length / n < KOverNThreshold)
{
PickKRandomIndexesIteratively(n, indexes, seed, distinct);
}
else
{
PickKRandomDistinctIndexesWithShuffle(n, indexes, seed);
}
}

/// <summary>
Expand All @@ -43,46 +47,39 @@ public static Span<int> PickKRandomIndexes(int n, int k, int seed, bool distinct
public static int PickRandomIndex(int n, int rand)
=> rand % n;

private static Span<int> PickKRandomIndexesIteratively(int n, int k, int seed, bool distinct)
private static void PickKRandomIndexesIteratively(int n, Span<int> indexes, int seed, bool distinct)
{
var random = new Random(seed);
var result = new int[k];
HashSet<int> pickedIndexes = default;
if (distinct) pickedIndexes = [];

var i = 0;
while (i < k)
while (i < indexes.Length)
{
var nextIdx = random.Next(n);
if (!distinct || pickedIndexes.Add(nextIdx))
{
result[i] = nextIdx;
indexes[i] = nextIdx;
i++;
}
}

return result;
}

private static Span<int> PickKRandomDistinctIndexesWithShuffle(int n, int k, int seed)
private static void PickKRandomDistinctIndexesWithShuffle(int n, Span<int> indexes, int seed)
{
const int StackallocThreshold = 256;

var random = new Random(seed);
var shuffledIndexes = new int[n];
for (var i = 0; i < n; i++)
{
shuffledIndexes[i] = i;
}
var shuffledIndexes = n <= StackallocThreshold ?
stackalloc int[StackallocThreshold].Slice(0, n) : new int[n];

// Fisher-Yates shuffle
for (var i = n - 1; i > 0; i--)
for (var i = 0; i < shuffledIndexes.Length; i++)
{
// Get a random index from 0 to i
var j = random.Next(i + 1);
// Swap shuffledIndexes[i] and shuffledIndexes[j]
(shuffledIndexes[i], shuffledIndexes[j]) = (shuffledIndexes[j], shuffledIndexes[i]);
shuffledIndexes[i] = i;
}

return new Span<int>(shuffledIndexes, 0, k);
random.Shuffle(shuffledIndexes);
shuffledIndexes.Slice(0, indexes.Length).CopyTo(indexes);
}
}
}
12 changes: 9 additions & 3 deletions libs/server/Objects/Hash/HashObjectImpl.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -132,11 +132,17 @@ private void HashRandomField(ref ObjectInput input, ref GarnetObjectStoreOutput
if (countParameter > 0 && countParameter > count)
countParameter = count;

var absCount = Math.Abs(countParameter);
var indexes = RandomUtils.PickKRandomIndexes(count, absCount, seed, countParameter > 0);
const int StackallocThreshold = 256;

var indexCount = Math.Abs(countParameter);

var indexes = indexCount <= StackallocThreshold ?
stackalloc int[StackallocThreshold].Slice(0, indexCount) : new int[indexCount];

RandomUtils.PickKRandomIndexes(count, indexes, seed, countParameter > 0);

// Write the size of the array reply
writer.WriteArrayLength(withValues && (respProtocolVersion == 2) ? absCount * 2 : absCount);
writer.WriteArrayLength(withValues && (respProtocolVersion == 2) ? indexCount * 2 : indexCount);

foreach (var index in indexes)
{
Expand Down
16 changes: 12 additions & 4 deletions libs/server/Objects/Set/SetObjectImpl.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -163,15 +163,18 @@ private void SetRandomMember(ref ObjectInput input, ref GarnetObjectStoreOutput

using var writer = new RespMemoryWriter(respProtocolVersion, ref output.SpanByteAndMemory);

Span<int> indexes = default;

if (count > 0)
{
// Return an array of distinct elements
var countParameter = count > Set.Count ? Set.Count : count;

// The order of fields in the reply is not truly random
indexes = RandomUtils.PickKRandomIndexes(Set.Count, countParameter, seed);
const int StackallocThreshold = 256;

var indexes = countParameter <= StackallocThreshold ?
stackalloc int[StackallocThreshold].Slice(0, countParameter) : new int[countParameter];

RandomUtils.PickKRandomIndexes(countParameter, indexes, seed);

// Write the size of the array reply
writer.WriteSetLength(countParameter);
Expand Down Expand Up @@ -203,9 +206,14 @@ private void SetRandomMember(ref ObjectInput input, ref GarnetObjectStoreOutput
else // count < 0
{
// Return an array with potentially duplicate elements
const int StackallocThreshold = 256;

var countParameter = Math.Abs(count);

indexes = RandomUtils.PickKRandomIndexes(Set.Count, countParameter, seed, false);
var indexes = countParameter <= StackallocThreshold ?
stackalloc int[StackallocThreshold].Slice(0, countParameter) : new int[countParameter];

RandomUtils.PickKRandomIndexes(Set.Count, indexes, seed, false);

if (Set.Count > 0)
{
Expand Down
9 changes: 8 additions & 1 deletion libs/server/Objects/SortedSet/SortedSetObjectImpl.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -615,7 +615,14 @@ private void SortedSetRandomMember(ref ObjectInput input, ref GarnetObjectStoreO
writer.WriteArrayLength(arrayLength);
}

var indexes = RandomUtils.PickKRandomIndexes(sortedSetCount, Math.Abs(count), seed, count > 0);
const int StackallocThreshold = 256;

var indexCount = Math.Abs(count);

var indexes = indexCount <= StackallocThreshold ?
stackalloc int[StackallocThreshold].Slice(0, indexCount) : new int[indexCount];

RandomUtils.PickKRandomIndexes(sortedSetCount, indexes, seed, count > 0);

foreach (var item in indexes)
{
Expand Down