Skip to content

Commit

Permalink
optimizations
Browse files Browse the repository at this point in the history
  • Loading branch information
@aizuon
aizuon committed Dec 22, 2023
1 parent bb9baf1 commit 0217aa2
Show file tree
Hide file tree
Showing 7 changed files with 129 additions and 90 deletions.
2 changes: 1 addition & 1 deletion tiny-dns/Packets/DNSHeader.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,7 +63,7 @@ public BinaryBuffer Serialize()
flags.Write(RA, 1);
flags.Write(Z, 3);
flags.Write(RCode, 4);
buffer.WriteRaw(flags.Buffer);
buffer.WriteRaw(flags.Buffer.AsSpan());

buffer.Write(Questions);
buffer.Write(AnswerRRs);
Expand Down
6 changes: 3 additions & 3 deletions tiny-dns/Packets/DNSQuery.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,12 +14,12 @@ public BinaryBuffer Serialize()

if (EDNSOption != null)
Header.AdditionalRRs = 1;
buffer.WriteRaw(Header.Serialize().Buffer);
buffer.WriteRaw(Header.Serialize().Buffer.AsSpan());

buffer.WriteRaw(Question.Serialize().Buffer);
buffer.WriteRaw(Question.Serialize().Buffer.AsSpan());

if (EDNSOption != null)
buffer.WriteRaw(EDNSOption.Serialize().Buffer);
buffer.WriteRaw(EDNSOption.Serialize().Buffer.AsSpan());

return buffer;
}
Expand Down
2 changes: 1 addition & 1 deletion tiny-dns/Packets/EDNSOption.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -64,7 +64,7 @@ public BinaryBuffer Serialize()
buffer.Write(rdLength);

foreach (var option in Options)
buffer.WriteRaw(option.Serialize().Buffer);
buffer.WriteRaw(option.Serialize().Buffer.AsSpan());

return buffer;
}
Expand Down
2 changes: 1 addition & 1 deletion tiny-dns/Packets/EDNSOptionData.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@ public BinaryBuffer Serialize()

buffer.Write((ushort)Data.Length);

buffer.WriteRaw(Data);
buffer.WriteRaw(Data.AsSpan());

return buffer;
}
Expand Down
2 changes: 1 addition & 1 deletion tiny-dns/RecursiveResolver.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -40,7 +40,7 @@ private static async ValueTask<IPAddress> ResolveRecursive(string qname, string

using var client = new UdpClient();
client.Connect(IPAddress.Parse(server), 53);
await client.SendAsync(req.Buffer, req.Buffer.Length);
await client.SendAsync(req.Buffer.AsMemory());
var res = await client.ReceiveAsync();

var buffer = new BinaryBuffer(res.Buffer);
Expand Down
176 changes: 111 additions & 65 deletions tiny-dns/Serialization/BinaryBuffer.cs
View file
Original file line number Diff line number Diff line change
@@ -1,69 +1,122 @@
using System.Numerics;
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;
using System.Text;
using ArgumentOutOfRangeException = System.ArgumentOutOfRangeException;

namespace TinyDNS.Serialization;

public class BinaryBuffer
{
private readonly object _mutex = new object();

private byte[] _buffer;

private uint _readOffset;

private uint _writeOffset;

public BinaryBuffer()
{
Buffer = Array.Empty<byte>();
_buffer = Array.Empty<byte>();
Capacity = 0;
Length = 0;
}

public BinaryBuffer(byte[] obj)
{
Buffer = obj;
WriteOffset = (uint)obj.Length;
_buffer = obj;
_writeOffset = (uint)obj.Length;
Capacity = (uint)obj.Length;
Length = (uint)obj.Length;
}

public byte[] Buffer { get; set; }
public ArraySegment<byte> Buffer
{
get
{
lock (_mutex)
{
return new ArraySegment<byte>(_buffer, 0, (int)Length);
}
}
}

public uint ReadOffset { get; set; }
public uint WriteOffset { get; set; }
public uint Length { get; private set; }

public unsafe void Write<T>(T obj) where T : unmanaged, IBinaryNumber<T>
public uint Capacity { get; private set; }

public uint ReadOffset
{
get => _readOffset;
set
{
if (value > Length)
throw new ArgumentOutOfRangeException();
_readOffset = value;
}
}

public uint WriteOffset
{
get => _writeOffset;
set
{
if (value > Length)
throw new ArgumentOutOfRangeException();
_writeOffset = value;
}
}

public void Write<T>(T obj) where T : unmanaged, IBinaryNumber<T>
{
lock (_mutex)
{
uint length = (uint)Unsafe.SizeOf<T>();
GrowIfNeeded(length);
fixed (byte* b = Buffer)
{
obj = Mem.ToBigEndian(obj);
int length = Unsafe.SizeOf<T>();
GrowIfNeeded((uint)length);

System.Buffer.MemoryCopy(&obj, b + WriteOffset, Buffer.Length - WriteOffset, length);
}
var bufferSpan = _buffer.AsSpan((int)_writeOffset, length);
obj = Mem.ToBigEndian(obj);
MemoryMarshal.Write(bufferSpan, in obj);

WriteOffset += length;
_writeOffset += (uint)length;
Length = Math.Max(Length, _writeOffset);
}
}

public void WriteRaw<T>(T[] obj) where T : unmanaged, IBinaryNumber<T>
public void WriteRaw<T>(Span<T> obj) where T : unmanaged, IBinaryNumber<T>
{
lock (_mutex)
{
uint length = (uint)obj.Length * (uint)Unsafe.SizeOf<T>();
int sizeOfT = Unsafe.SizeOf<T>();
uint length = (uint)obj.Length * (uint)sizeOfT;
GrowIfNeeded(length);

foreach (var o in obj)
Write(o);
var targetSpan = new Span<byte>(_buffer, (int)_writeOffset, (int)length);

for (int i = 0; i < obj.Length; i++)
{
var value = Mem.ToBigEndian(obj[i]);
MemoryMarshal.Write(targetSpan.Slice(i * sizeOfT, sizeOfT), in value);
}

_writeOffset += length;
Length = Math.Max(Length, _writeOffset);
}
}

public void Write(string obj)
public void WriteString(string obj)
{
lock (_mutex)
{
byte[] ascii = Encoding.ASCII.GetBytes(obj);
Write((byte)ascii.Length);
int byteCount = Encoding.ASCII.GetByteCount(obj);
Write((byte)byteCount);

GrowIfNeeded((uint)ascii.Length);
GrowIfNeeded((uint)byteCount);

foreach (byte o in ascii)
Write(o);
Encoding.ASCII.GetBytes(obj, 0, obj.Length, _buffer, (int)_writeOffset);
_writeOffset += (uint)byteCount;
Length = Math.Max(Length, _writeOffset);
}
}

Expand All @@ -73,53 +126,51 @@ public void WriteDomainName(string obj)
{
string[] labels = obj.Split('.');
foreach (string label in labels)
Write(label);
WriteString(label);
Write((byte)0);
}
}

public unsafe T Read<T>() where T : unmanaged, IBinaryNumber<T>
public T Read<T>() where T : unmanaged, IBinaryNumber<T>
{
lock (_mutex)
{
uint length = (uint)Unsafe.SizeOf<T>();
int length = Unsafe.SizeOf<T>();

uint finalOffset = ReadOffset + length;
if (Buffer.Length < finalOffset)
if (_buffer.Length < _readOffset + length)
throw new ArgumentOutOfRangeException();

var obj = T.Zero;
var o = &obj;
fixed (byte* b = Buffer)
{
byte* p = (byte*)o;

System.Buffer.MemoryCopy(b + ReadOffset, p, length, length);
var bufferSpan = _buffer.AsSpan((int)_readOffset, length);
var value = MemoryMarshal.Read<T>(bufferSpan);
value = Mem.ToBigEndian(value);

obj = Mem.ToBigEndian(obj);
}

ReadOffset = finalOffset;
_readOffset += (uint)length;

return obj;
return value;
}
}

public T[] ReadRaw<T>(uint count) where T : unmanaged, IBinaryNumber<T>
{
lock (_mutex)
{
uint length = count * (uint)Unsafe.SizeOf<T>();
int sizeOfT = Unsafe.SizeOf<T>();
uint length = count * (uint)sizeOfT;

uint finalOffset = ReadOffset + length;
if (Buffer.Length < finalOffset)
if (_buffer.Length - _readOffset < length)
throw new ArgumentOutOfRangeException();

var obj = new T[count];
for (uint i = 0; i < count; i++)
obj[i] = Read<T>();
var result = new T[count];
var byteSpan = new Span<byte>(_buffer, (int)_readOffset, (int)length);

return obj;
for (int i = 0; i < count; i++)
{
var value = MemoryMarshal.Read<T>(byteSpan[(i * sizeOfT)..]);
result[i] = Mem.ToBigEndian(value);
}

_readOffset += length;
return result;
}
}

Expand All @@ -129,17 +180,11 @@ public string ReadString()
{
byte size = Read<byte>();

uint length = (uint)(size * sizeof(byte));

uint finalOffset = ReadOffset + length;
if (Buffer.Length < finalOffset)
if (_buffer.Length < _readOffset + size)
throw new ArgumentOutOfRangeException();

byte[] ascii = new byte[length];
for (uint i = 0; i < size; i++)
ascii[i] = Read<byte>();

string obj = Encoding.ASCII.GetString(ascii);
string obj = Encoding.ASCII.GetString(_buffer, (int)_readOffset, size);
_readOffset += size;

return obj;
}
Expand Down Expand Up @@ -188,14 +233,15 @@ public string ReadDomainName()

private void GrowIfNeeded(uint writeLength)
{
uint finalLength = WriteOffset + writeLength;
bool resizeNeeded = Buffer.Length <= finalLength;

if (resizeNeeded)
uint requiredLength = WriteOffset + writeLength;
if (Capacity < requiredLength)
{
byte[] tmp = Buffer;
Array.Resize(ref tmp, (int)finalLength);
Buffer = tmp;
uint newCapacity = Math.Max(Capacity * 2, requiredLength);
byte[] tmp = _buffer;
Array.Resize(ref tmp, (int)newCapacity);
_buffer = tmp;
Length = requiredLength;
Capacity = newCapacity;
}
}
}
29 changes: 11 additions & 18 deletions tiny-dns/Serialization/Mem.cs
View file
Original file line number Diff line number Diff line change
@@ -1,30 +1,23 @@
using System.Runtime.CompilerServices;
using System.Runtime.InteropServices;

namespace TinyDNS.Serialization;

public static class Mem
{
public static unsafe T ToBigEndian<T>(T obj) where T : unmanaged
{
if (BitConverter.IsLittleEndian)
{
uint length = (uint)Unsafe.SizeOf<T>();
if (!BitConverter.IsLittleEndian)
return obj;

var o = &obj;
byte* p = (byte*)o;
if (BitConverter.IsLittleEndian)
{
byte* pStart = p;
byte* pEnd = p + length - 1;
for (int i = 0; i < length / 2; i++)
{
byte temp = *pStart;
*pStart++ = *pEnd;
*pEnd-- = temp;
}
}
}
int size = Unsafe.SizeOf<T>();
Span<byte> bytes = stackalloc byte[size];
ref byte objRef = ref Unsafe.As<T, byte>(ref obj);

return obj;
Unsafe.CopyBlockUnaligned(ref MemoryMarshal.GetReference(bytes), ref objRef, (uint)size);

bytes.Reverse();

return MemoryMarshal.Read<T>(bytes);
}
}

0 comments on commit 0217aa2

Please sign in to comment.