void DisposeAssembler(StreamMessageAssembler assembler)
{
assembler.Dispose();
_assemblers.Remove(assembler.MessageId);
}
public async Task <Message> ReceiveMessage(CancellationToken ct)
{
var buffer = ArrayPool <byte> .Shared.Rent(_receiveBufferSize);
try
{
while (true)
{
ct.ThrowIfCancellationRequested();
var read = 0;
var partialMsgSize = -1;
var msgSize = 0;
while (partialMsgSize == -1 || read < partialMsgSize)
{
var len = partialMsgSize == -1 ? 8 : partialMsgSize - read;
var segment = new ArraySegment <byte>(buffer, read, len);
read += await _transportChannel.ReceiveAsync(segment, ct);
if (partialMsgSize == -1 && read >= 8)
{
partialMsgSize = BitConverter.ToInt32(buffer, _partialMsgSizeIndex);
msgSize = BitConverter.ToInt32(buffer, _msgSizeIndex);
if (partialMsgSize < _minPartialMsgSize)
{
throw new Exception($"Received message size {partialMsgSize} while minimum is {_minPartialMsgSize}");
}
if (partialMsgSize > _receiveBufferSize)
{
var newBuffer = ArrayPool <byte> .Shared.Rent(partialMsgSize);
Buffer.BlockCopy(buffer, 0, newBuffer, 0, read);
ArrayPool <byte> .Shared.Return(buffer);
buffer = newBuffer;
}
if (msgSize > _maxMessageSize)
{
throw new Exception($"Received message size {partialMsgSize} while maximum is {_maxMessageSize}");
}
}
ct.ThrowIfCancellationRequested();
}
var msgId = new Guid(BufferHelper.GetSubArray(buffer, _idIndex, 16));
if (msgSize == partialMsgSize - _msgSizeIndex)
{
var msgType = (MessageType)buffer[_partialMsgTypeIndex];
byte[] payload = null;
var payloadLength = partialMsgSize - _partialPayloadIndex;
if (payloadLength > 0)
{
payload = BufferHelper.GetSubArray(buffer, _partialPayloadIndex, payloadLength);
}
return(new Message(msgType, msgId, payload));
}
else
{
var nextIndex = _msgSizeIndex;
if (!_assemblers.TryGetValue(msgId, out var assembler))
{
assembler = new StreamMessageAssembler(msgId, msgSize);
_assemblers.Add(msgId, assembler);
}
else
{
nextIndex = _partialPayloadIndex;
}
var nextSegment = new ArraySegment <byte>(buffer, nextIndex, partialMsgSize - nextIndex);
if (assembler.AddPartialMessage(nextSegment))
{
var msgBuffer = assembler.GetMessage();
var msgType = (MessageType)msgBuffer[_msgTypeIndex];
byte[] payload = null;
var payloadLength = msgSize - _payloadIndex;
if (payloadLength > 0)
{
payload = BufferHelper.GetSubArray(msgBuffer, _payloadIndex, payloadLength);
}
var msg = new Message(msgType, msgId, payload);
DisposeAssembler(assembler);
return(msg);
}
}
}
}
finally
{
ArrayPool <byte> .Shared.Return(buffer);
}
}
