Skip to content

perf: get rid of MemoryStream in KeyRingBasedDataProtector #59322

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
merged 7 commits into from
Jan 8, 2025
Merged

perf: get rid of MemoryStream in KeyRingBasedDataProtector #59322

Changes from 2 commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
c686038
get rid of memoryStream usage
DeagleGross Dec 4, 2024
8ac940f
dont console writeline :)
DeagleGross Dec 4, 2024
4b5ca1a
add tests and address PR comments
DeagleGross Dec 5, 2024
a046b3f
standartize API + use stackalloc for small arrays
DeagleGross Dec 5, 2024
f233b4e
Merge branch 'main' into dmkorolev/dataprotection/lin-perf-1
DeagleGross Jan 7, 2025
51dec55
address PR comments
DeagleGross Jan 7, 2025
8212333
move utils to Shared\Encoding
DeagleGross Jan 8, 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
147 changes: 116 additions & 31 deletions src/DataProtection/DataProtection/src/KeyManagement/KeyRingBasedDataProtector.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,8 @@
// The .NET Foundation licenses this file to you under the MIT license.

using System;
using System.Buffers.Binary;
using System.Buffers;
using System.Collections.Generic;
using System.Diagnostics;
using System.Diagnostics.CodeAnalysis;
Expand All @@ -14,6 +16,7 @@
using Microsoft.AspNetCore.DataProtection.KeyManagement.Internal;
using Microsoft.AspNetCore.Shared;
using Microsoft.Extensions.Logging;
using System.Buffers.Text;

namespace Microsoft.AspNetCore.DataProtection.KeyManagement;

Expand Down Expand Up @@ -317,37 +320,6 @@ private struct AdditionalAuthenticatedDataTemplate
{
private byte[] _aadTemplate;

public AdditionalAuthenticatedDataTemplate(IEnumerable<string> purposes)
{
const int MEMORYSTREAM_DEFAULT_CAPACITY = 0x100; // matches MemoryStream.EnsureCapacity
var ms = new MemoryStream(MEMORYSTREAM_DEFAULT_CAPACITY);

// additionalAuthenticatedData := { magicHeader (32-bit) || keyId || purposeCount (32-bit) || (purpose)* }
// purpose := { utf8ByteCount (7-bit encoded) || utf8Text }

using (var writer = new PurposeBinaryWriter(ms))
{
writer.WriteBigEndian(MAGIC_HEADER_V0);
Debug.Assert(ms.Position == sizeof(uint));
var posPurposeCount = writer.Seek(sizeof(Guid), SeekOrigin.Current); // skip over where the key id will be stored; we'll fill it in later
writer.Seek(sizeof(uint), SeekOrigin.Current); // skip over where the purposeCount will be stored; we'll fill it in later

uint purposeCount = 0;
foreach (string purpose in purposes)
{
Debug.Assert(purpose != null);
writer.Write(purpose); // prepends length as a 7-bit encoded integer
purposeCount++;
}

// Once we have written all the purposes, go back and fill in 'purposeCount'
writer.Seek(checked((int)posPurposeCount), SeekOrigin.Begin);
writer.WriteBigEndian(purposeCount);
}

_aadTemplate = ms.ToArray();
}

public byte[] GetAadForKey(Guid keyId, bool isProtecting)
{
// Multiple threads might be trying to read and write the _aadTemplate field
Expand Down Expand Up @@ -381,6 +353,118 @@ public byte[] GetAadForKey(Guid keyId, bool isProtecting)
}
}

#if NET10_0_OR_GREATER
public AdditionalAuthenticatedDataTemplate(string[] purposes)
{
// additionalAuthenticatedData := { magicHeader (32-bit) || keyId || purposeCount (32-bit) || (purpose)* }
// purpose := { utf8ByteCount (7-bit encoded) || utf8Text }

var keySize = sizeof(Guid);
int totalPurposeLen = 4 + keySize + 4;

var purposeLengthsPool = ArrayPool<int>.Shared.Rent(purposes.Length);
Copy link
Member

@sebastienros sebastienros Dec 4, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Did you measure if it was worth renting an array only to prevent an extra call to SecureUtf8Encoding.GetByteCount and Measure7BitEncodedUIntLength

Copy link
Contributor

@mgravell mgravell Dec 5, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

You mean just pay the cost twice,? Yeah, should probably measure it.

Side thought: I'm guessing we don't actually expect many purposes. Another thing we could do here, if you're concerned about ArrayPool overhead, is use the stack - something like:

int[]? lease = null;
Span<int> lengths = purposes.Length <= 32 ? stackalloc int[purposes.Length] : (lease = ArrayPool<int>.Shared.Rent(purposes.Length)).AsSpan(0, purposes.Length);
...
if (lease is not null)
{
    ArrayPool<int>.Shared.Return(lease);
}

Copy link
Member Author

@DeagleGross DeagleGross Dec 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Benchmark code

Method Mean Error StdDev Gen0 Gen1 Gen2 Allocated
MemoryStream 141.4 ns 2.62 ns 3.31 ns 0.0067 - - 512 B
Manual 101.8 ns 0.98 ns 0.92 ns 0.0023 0.0002 0.0002 320 B
Manual_CalculateByteCountEveryTime 124.7 ns 1.62 ns 1.52 ns 0.0017 - - 128 B
Manual_StackAllocForSmallPurposeArrays 122.4 ns 0.70 ns 0.54 ns 0.0014 - - 128 B

It seems like that last option Marc suggested is the most balanced - we indeed dont expect many purposes (I think), and this solution does not contest for the ArrayPool rent.

What do you think?

Copy link
Contributor

@mgravell mgravell Dec 5, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

are the last two doing the same work as Manual? it seems odd that the perf is worse than Manual - we're doing less work (by not touching the pool at all) - why is it taking longer? question: can you try this with [SkipLocalsInit] on the method? or the entire type for fairness (this is applied globally in aspnetcore, I believe, but I don't know whether it'll apply to your benchmark which looks to be custom).

I also don't know why the memory is less if it is doing the same work - the only non-amortized allocation here is the result byte[]; so I assume this is a different data test than Manual?

Copy link
Member

@BrennanConroy BrennanConroy Dec 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you try this with [SkipLocalsInit] on the method? or the entire type for fairness (this is applied globally in aspnetcore, I believe

Don't think it is. #26586

Copy link
Member Author

@DeagleGross DeagleGross Dec 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

we have talked with Marc: he launched the benchmarks on his machine and there is no crazy allocs which are confusing. I think we should merge with the current state of code (stackalloc for under 32 dynamic length of purposes)

Copy link
Contributor

@mgravell mgravell Dec 5, 2024

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

can you try this with [SkipLocalsInit] on the method? or the entire type for fairness (this is applied globally in aspnetcore, I believe

Don't think it is. #26586

We should revisit this, IMHO.

for (int i = 0; i < purposes.Length; i++)
{
string purpose = purposes[i];

int purposeLength = EncodingUtil.SecureUtf8Encoding.GetByteCount(purpose);
purposeLengthsPool[i] = purposeLength;

var encoded7BitUIntLength = Measure7BitEncodedUIntLength((uint)purposeLength);
totalPurposeLen += purposeLength /* length of actual string */ + encoded7BitUIntLength /* length of 'string length' 7-bit encoded int */;
}

byte[] targetArr = new byte[totalPurposeLen];
var targetSpan = targetArr.AsSpan();

// index 0: magic header
BinaryPrimitives.WriteUInt32BigEndian(targetSpan.Slice(0), MAGIC_HEADER_V0);
// index 4: key (skipped for now, will be populated in `GetAadForKey()`)
// index 4 + keySize: purposeCount
BinaryPrimitives.WriteInt32BigEndian(targetSpan.Slice(4 + keySize), purposes.Length);

int index = 4 + keySize + 4; // starting from first purpose
for (int i = 0; i < purposes.Length; i++)
{
string purpose = purposes[i];

// writing `utf8ByteCount (7-bit encoded integer) || utf8Text`
// we have already calculated the lengths of the purpose strings, so just get it from the pool
index += Write7BitEncodedInt(purposeLengthsPool[i], targetSpan.Slice(index));
index += EncodingUtil.SecureUtf8Encoding.GetBytes(purpose.AsSpan(), targetSpan.Slice(index));
}

ArrayPool<int>.Shared.Return(purposeLengthsPool);
Debug.Assert(index == targetArr.Length);

_aadTemplate = targetArr;
}

private static int Measure7BitEncodedUIntLength(uint value)
{
return ((31 - System.Numerics.BitOperations.LeadingZeroCount(value | 1)) / 7) + 1;

// does the same as the following code:
// int count = 1;
// while ((value >>= 7) != 0)
// {
// count++;
// }
// return count;
}

private static int Write7BitEncodedInt(int value, Span<byte> target)
{
uint uValue = (uint)value;

// Write out an int 7 bits at a time. The high bit of the byte,
// when on, tells reader to continue reading more bytes.
//
// Using the constants 0x7F and ~0x7F below offers smaller
// codegen than using the constant 0x80.

int index = 0;
while (uValue > 0x7Fu)
{
target[index++] = (byte)(uValue | ~0x7Fu);
uValue >>= 7;
}

target[index++] = (byte)uValue;
return index;
}
#else
public AdditionalAuthenticatedDataTemplate(IEnumerable<string> purposes)
{
const int MEMORYSTREAM_DEFAULT_CAPACITY = 0x100; // matches MemoryStream.EnsureCapacity
var ms = new MemoryStream(MEMORYSTREAM_DEFAULT_CAPACITY);

// additionalAuthenticatedData := { magicHeader (32-bit) || keyId || purposeCount (32-bit) || (purpose)* }
// purpose := { utf8ByteCount (7-bit encoded) || utf8Text }

using (var writer = new PurposeBinaryWriter(ms))
{
writer.WriteBigEndian(MAGIC_HEADER_V0);
Debug.Assert(ms.Position == sizeof(uint));
var posPurposeCount = writer.Seek(sizeof(Guid), SeekOrigin.Current); // skip over where the key id will be stored; we'll fill it in later
writer.Seek(sizeof(uint), SeekOrigin.Current); // skip over where the purposeCount will be stored; we'll fill it in later

uint purposeCount = 0;
foreach (string purpose in purposes)
{
Debug.Assert(purpose != null);
writer.Write(purpose); // prepends length as a 7-bit encoded integer
purposeCount++;
}

// Once we have written all the purposes, go back and fill in 'purposeCount'
writer.Seek(checked((int)posPurposeCount), SeekOrigin.Begin);
writer.WriteBigEndian(purposeCount);
}

_aadTemplate = ms.ToArray();
}

private sealed class PurposeBinaryWriter : BinaryWriter
{
public PurposeBinaryWriter(MemoryStream stream) : base(stream, EncodingUtil.SecureUtf8Encoding, leaveOpen: true) { }
Expand All @@ -395,6 +479,7 @@ public void WriteBigEndian(uint value)
outStream.WriteByte((byte)(value));
}
}
#endif
}

private enum UnprotectStatus
Expand Down
Loading