Skip to content

Commit

Permalink
Rewrote send to use preallocated/pooled buffers, too
Browse files Browse the repository at this point in the history
  • Loading branch information
@MichaCo
MichaCo committed Jan 7, 2017
1 parent a5abaa9 commit 04478f5
Show file tree
Hide file tree
Showing 15 changed files with 161 additions and 151 deletions.
2 changes: 1 addition & 1 deletion samples/MiniDig/Properties/launchSettings.json
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,7 +2,7 @@
"profiles": {
"DigApp": {
"commandName": "Project",
"commandLineArgs": "perf -s 192.168.178.21 localhost -c 5 -r 10"
"commandLineArgs": "-s 8.8.4.4 microsoft.com any"
},
"MiniDig": {
"commandName": "Project"
Expand Down
20 changes: 14 additions & 6 deletions src/DnsClient/DnsDatagramWriter.cs
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
using System;
using System.Diagnostics;
using System.Linq;
using System.Net;
using System.Text;
Expand All @@ -8,28 +9,35 @@ namespace DnsClient
internal class DnsDatagramWriter : IDisposable
{
// queries can only be 255 octets + some header bytes, so that size is pretty safe...
private const int MaxBufferSize = 1024;
public const int BufferSize = 1024;

private const byte DotByte = 46;

private readonly PooledBytes _pooledBytes;

private ArraySegment<byte> _buffer;

public byte[] Data
public ArraySegment<byte> Data
{
get
{
return new ArraySegment<byte>(_buffer.Array, 0, Index).ToArray();
return new ArraySegment<byte>(_buffer.Array, 0, Index);
}
}

public int Index { get; set; }

public DnsDatagramWriter()
{
_pooledBytes = new PooledBytes(MaxBufferSize);
_buffer = new ArraySegment<byte>(_pooledBytes.Buffer, 0, MaxBufferSize);
_pooledBytes = new PooledBytes(BufferSize);
_buffer = new ArraySegment<byte>(_pooledBytes.Buffer, 0, BufferSize);
}

public DnsDatagramWriter(ArraySegment<byte> useBuffer)
{
Debug.Assert(useBuffer.Count >= BufferSize);

_buffer = useBuffer;
}

public void WriteHostName(string queryName)
Expand Down Expand Up @@ -91,7 +99,7 @@ protected virtual void Dispose(bool disposing)
{
if (disposing)
{
_pooledBytes.Dispose();
_pooledBytes?.Dispose();
}
}

Expand Down
105 changes: 49 additions & 56 deletions src/DnsClient/DnsMessageHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -14,67 +14,60 @@ internal abstract class DnsMessageHandler

public abstract bool IsTransientException<T>(T exception) where T : Exception;

public virtual byte[] GetRequestData(DnsRequestMessage request)
public virtual void GetRequestData(DnsRequestMessage request, DnsDatagramWriter writer)
{
var question = request.Question;

/*
1 1 1 1 1 1
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| ID |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
|QR| Opcode |AA|TC|RD|RA| Z | RCODE |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| QDCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| ANCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| NSCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| ARCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
* */
1 1 1 1 1 1
0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| ID |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
|QR| Opcode |AA|TC|RD|RA| Z | RCODE |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| QDCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| ANCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| NSCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
| ARCOUNT |
+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+--+
* */
// 4 more bytes for the type and class
using (var writer = new DnsDatagramWriter())
{
writer.WriteInt16NetworkOrder((short)request.Header.Id);
writer.WriteUInt16NetworkOrder(request.Header.RawFlags);
writer.WriteInt16NetworkOrder(1); // we support single question only... (as most DNS servers anyways).
writer.WriteInt16NetworkOrder(0);
writer.WriteInt16NetworkOrder(0);
writer.WriteInt16NetworkOrder(1); // one additional for the Opt record.

writer.WriteHostName(question.QueryName);
//writer.WriteBytes(questionData.Array, questionData.Count);
writer.WriteUInt16NetworkOrder((ushort)question.QuestionType);
writer.WriteUInt16NetworkOrder((ushort)question.QuestionClass);

/*
+------------+--------------+------------------------------+
| Field Name | Field Type | Description |
+------------+--------------+------------------------------+
| NAME | domain name | MUST be 0 (root domain) |
| TYPE | u_int16_t | OPT (41) |
| CLASS | u_int16_t | requestor's UDP payload size |
| TTL | u_int32_t | extended RCODE and flags |
| RDLEN | u_int16_t | length of all RDATA |
| RDATA | octet stream | {attribute,value} pairs |
+------------+--------------+------------------------------+
* */
var opt = new OptRecord();
//var nameBytes = opt.DomainName.GetBytes();
writer.WriteHostName("");
//writer.WriteBytes(nameBytes.Array, nameBytes.Count);
writer.WriteUInt16NetworkOrder((ushort)opt.RecordType);
writer.WriteUInt16NetworkOrder((ushort)opt.RecordClass);
writer.WriteUInt32NetworkOrder((ushort)opt.TimeToLive);
writer.WriteUInt16NetworkOrder(0);

return writer.Data;

// dispose the writer here to return puled byte array... otherwise we don't know when to dispose and might be risky
}

writer.WriteInt16NetworkOrder((short)request.Header.Id);
writer.WriteUInt16NetworkOrder(request.Header.RawFlags);
writer.WriteInt16NetworkOrder(1); // we support single question only... (as most DNS servers anyways).
writer.WriteInt16NetworkOrder(0);
writer.WriteInt16NetworkOrder(0);
writer.WriteInt16NetworkOrder(1); // one additional for the Opt record.

writer.WriteHostName(question.QueryName);
writer.WriteUInt16NetworkOrder((ushort)question.QuestionType);
writer.WriteUInt16NetworkOrder((ushort)question.QuestionClass);

/*
+------------+--------------+------------------------------+
| Field Name | Field Type | Description |
+------------+--------------+------------------------------+
| NAME | domain name | MUST be 0 (root domain) |
| TYPE | u_int16_t | OPT (41) |
| CLASS | u_int16_t | requestor's UDP payload size |
| TTL | u_int32_t | extended RCODE and flags |
| RDLEN | u_int16_t | length of all RDATA |
| RDATA | octet stream | {attribute,value} pairs |
+------------+--------------+------------------------------+
* */

var opt = new OptRecord();

writer.WriteHostName("");
writer.WriteUInt16NetworkOrder((ushort)opt.RecordType);
writer.WriteUInt16NetworkOrder((ushort)opt.RecordClass);
writer.WriteUInt32NetworkOrder((ushort)opt.TimeToLive);
writer.WriteUInt16NetworkOrder(0);
}

public virtual DnsResponseMessage GetResponseMessage(ArraySegment<byte> responseData)
Expand Down
2 changes: 1 addition & 1 deletion src/DnsClient/DnsName.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -447,7 +447,7 @@ public byte[] GetBytes()
using (var writer = new DnsDatagramWriter())
{
WriteBytes(writer);
return writer.Data;
return writer.Data.ToArray();
}
}

Expand Down
30 changes: 16 additions & 14 deletions src/DnsClient/DnsTcpMessageHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ public override bool IsTransientException<T>(T exception)

public override DnsResponseMessage Query(IPEndPoint endpoint, DnsRequestMessage request)
{
throw new NotImplementedException();
return QueryAsync(endpoint, request, CancellationToken.None).Result;
}

public override async Task<DnsResponseMessage> QueryAsync(
Expand All @@ -32,21 +32,23 @@ public override async Task<DnsResponseMessage> QueryAsync(
await client.ConnectAsync(server.Address, server.Port).ConfigureAwait(false);
using (var stream = client.GetStream())
{
var data = GetRequestData(request);
int dataLength = data.Length;

//var sendLength = new byte[2];
//sendLength[0] = (byte)((data.Length >> 8) & 0xff);
//sendLength[1] = (byte)(data.Length & 0xff);
// use a pooled buffer to writer the data + the length of the data later into the frist two bytes
using (var memory = new PooledBytes(DnsDatagramWriter.BufferSize + 2))
using (var writer = new DnsDatagramWriter(new ArraySegment<byte>(memory.Buffer, 2, memory.Buffer.Length - 2)))
{
GetRequestData(request, writer);
int dataLength = writer.Index;
memory.Buffer[0] = (byte)((dataLength >> 8) & 0xff);
memory.Buffer[1] = (byte)(dataLength & 0xff);

var sendData = new byte[dataLength + 2];
sendData[0] = (byte)((dataLength >> 8) & 0xff);
sendData[1] = (byte)(dataLength & 0xff);
Array.Copy(data, 0, sendData, 2, dataLength);

await stream.WriteAsync(sendData, 0, sendData.Length, cancellationToken).ConfigureAwait(false);
//var sendData = new byte[dataLength + 2];
//sendData[0] = (byte)((dataLength >> 8) & 0xff);
//sendData[1] = (byte)(dataLength & 0xff);
//Array.Copy(data, 0, sendData, 2, dataLength);

await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
await stream.WriteAsync(memory.Buffer, 0, dataLength + 2, cancellationToken).ConfigureAwait(false);
await stream.FlushAsync(cancellationToken).ConfigureAwait(false);
}

int length = stream.ReadByte() << 8 | stream.ReadByte();
if (length <= 0)
Expand Down
87 changes: 51 additions & 36 deletions src/DnsClient/DnsUdpMessageHandler.cs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -10,20 +10,13 @@ namespace DnsClient
{
internal class DnsUdpMessageHandler : DnsMessageHandler
{
private const int MaxSize = 4096 * 4;
private const int MaxSize = 4096;
private static ConcurrentQueue<UdpClient> _clients = new ConcurrentQueue<UdpClient>();
private readonly bool _enableClientQueue;

public DnsUdpMessageHandler(bool enableClientQueue)
{
_enableClientQueue = enableClientQueue;
if (_enableClientQueue && _clients.Count == 0)
{
for (var i = 0; i < 10; i++)
{
_clients.Enqueue(new UdpClient());
}
}
}

public override bool IsTransientException<T>(T exception)
Expand All @@ -39,18 +32,19 @@ public override DnsResponseMessage Query(
UdpClient udpClient = GetNextUdpClient();
try
{
var data = GetRequestData(request);
udpClient.Client.SendTo(data, server);

using (var memory = new PooledBytes(MaxSize))
using (var writer = new DnsDatagramWriter())
{
var received = udpClient.Client.Receive(memory.Buffer, 0, 4096, SocketFlags.None);
while (udpClient.Available > 0)
{
received = udpClient.Client.Receive(memory.Buffer, received, 4096, SocketFlags.None);
}
GetRequestData(request, writer);
udpClient.Client.SendTo(writer.Data.Array, writer.Data.Offset, writer.Data.Count, SocketFlags.None, server);
}

var readSize = udpClient.Available > MaxSize ? udpClient.Available : MaxSize;

var response = GetResponseMessage(new ArraySegment<byte>(memory.Buffer, 0, memory.Buffer.Length));
using (var memory = new PooledBytes(readSize))
{
var received = udpClient.Client.Receive(memory.Buffer, 0, readSize, SocketFlags.None);

var response = GetResponseMessage(new ArraySegment<byte>(memory.Buffer, 0, received));
if (request.Header.Id != response.Header.Id)
{
throw new DnsResponseException("Header id missmatch.");
Expand All @@ -68,16 +62,19 @@ public override DnsResponseMessage Query(
{
if (!_enableClientQueue)
{
try
{
#if PORTABLE
udpClient.Dispose();
udpClient.Dispose();
#else
udpClient.Close();
udpClient.Close();
#endif
}
catch { }
}
}
}

////private static readonly byte[] fixData = new byte[] { 25, 158, 133, 0, 0, 1, 0, 1, 0, 0, 0, 1, 9, 108, 111, 99, 97, 108, 104, 111, 115, 116, 0, 0, 1, 0, 1, 192, 12, 0, 1, 0, 1, 0, 9, 58, 128, 0, 4, 127, 0, 0, 1, 0, 0, 41, 16, 0, 0, 0, 0, 0, 0, 0 };
public override async Task<DnsResponseMessage> QueryAsync(
IPEndPoint server,
DnsRequestMessage request,
Expand All @@ -88,34 +85,52 @@ public override async Task<DnsResponseMessage> QueryAsync(
UdpClient udpClient = GetNextUdpClient();
try
{
var data = GetRequestData(request);
await udpClient.SendAsync(data, data.Length, server).ConfigureAwait(false);

var result = await udpClient.ReceiveAsync().ConfigureAwait(false);

var response = GetResponseMessage(new ArraySegment<byte>(result.Buffer, 0, result.Buffer.Length));

if (request.Header.Id != response.Header.Id)
using (var writer = new DnsDatagramWriter())
{
throw new DnsResponseException("Header id missmatch.");
GetRequestData(request, writer);
await udpClient.SendAsync(writer.Data.Array, writer.Data.Count, server).ConfigureAwait(false);
}

if (_enableClientQueue)
var readSize = udpClient.Available > MaxSize ? udpClient.Available : MaxSize;

using (var memory = new PooledBytes(readSize))
{
_clients.Enqueue(udpClient);
}
#if PORTABLE
int received = await udpClient.Client.ReceiveAsync(new ArraySegment<byte>(memory.Buffer), SocketFlags.None).ConfigureAwait(false);

var response = GetResponseMessage(new ArraySegment<byte>(memory.Buffer, 0, received));

#else
var result = await udpClient.ReceiveAsync().ConfigureAwait(false);

var response = GetResponseMessage(new ArraySegment<byte>(result.Buffer, 0, result.Buffer.Length));
#endif
if (request.Header.Id != response.Header.Id)
{
throw new DnsResponseException("Header id missmatch.");
}

if (_enableClientQueue)
{
_clients.Enqueue(udpClient);
}

return response;
return response;
}
}
finally
{
if (!_enableClientQueue)
{
try
{
#if PORTABLE
udpClient.Dispose();
udpClient.Dispose();
#else
udpClient.Close();
udpClient.Close();
#endif
}
catch { }
}
}
}
Expand Down
Loading

0 comments on commit 04478f5

Please sign in to comment.