private void ReceiveCallback(Socket socket, SocketAsyncEventArgs socketAsyncEventArgs)
{
// Check for error
if (socketAsyncEventArgs.SocketError != SocketError.Success)
{
(socketAsyncEventArgs.UserToken as PooledMessage)?.Dispose();
HandleCommunicationError(socket, new Exception("Receive failed, error = " + socketAsyncEventArgs.SocketError));
}
// Get the message state
int bytesRead = socketAsyncEventArgs.BytesTransferred;
// Read the data
if (bytesRead > 0)
{
// Add to receive queue
BlockingQueue <SocketAsyncEventArgs> receiveBufferQueue = null;
_currentlyConnectedClientsReceiveQueuesLock.EnterReadLock();
try
{
if (!_currentlyConnectedClientsReceiveQueues.TryGetValue(socket, out receiveBufferQueue))
{
// Peace out!
return;
}
}
finally
{
_currentlyConnectedClientsReceiveQueuesLock.ExitReadLock();
}
receiveBufferQueue.Enqueue(socketAsyncEventArgs);
}
else
{
// 0 bytes means disconnect
(socketAsyncEventArgs.UserToken as PooledMessage)?.Dispose();
HandleCommunicationError(socket, new Exception("Received 0 bytes (graceful disconnect)"));
_socketAsyncEventArgsReceivePool.Push(socketAsyncEventArgs);
return;
}
socketAsyncEventArgs = _socketAsyncEventArgsReceivePool.Pop();
// Post a receive to the socket as the client will be continuously receiving messages to be pushed to the queue
TryUnsafeSocketOperation(socket, SocketAsyncOperation.Receive, socketAsyncEventArgs);
}
/// <summary>
/// The constructor.
/// </summary>
/// <param name="socketFunc">The function that creates a new socket. Use this to specify your socket constructor and initialize settings.</param>
/// <param name="messageBufferSize">The message buffer size to use for send/receive.</param>
/// <param name="communicationTimeout">The communication timeout, in milliseconds.</param>
/// <param name="maxMessageSize">The maximum message size.</param>
/// <param name="useNagleAlgorithm">Whether or not to use the Nagle algorithm.</param>
public SimplSocketClient(Func <Socket> socketFunc, int messageBufferSize = 65536, int communicationTimeout = 10000, int maxMessageSize = 10 * 1024 * 1024, bool useNagleAlgorithm = false)
{
// Sanitize
if (socketFunc == null)
{
throw new ArgumentNullException("socketFunc");
}
if (messageBufferSize < 512)
{
throw new ArgumentException("must be >= 512", "messageBufferSize");
}
if (communicationTimeout < 5000)
{
throw new ArgumentException("must be >= 5000", "communicationTimeout");
}
if (maxMessageSize < 1024)
{
throw new ArgumentException("must be >= 1024", "maxMessageSize");
}
_socketFunc = socketFunc;
_messageBufferSize = messageBufferSize;
_communicationTimeout = communicationTimeout;
_maxMessageSize = maxMessageSize;
_useNagleAlgorithm = useNagleAlgorithm;
_sendBufferQueue = new BlockingQueue <SocketAsyncEventArgs>();
_receiveBufferQueue = new BlockingQueue <SocketAsyncEventArgs>();
// Initialize the client multiplexer
_clientMultiplexer = new Dictionary <int, MultiplexerData>(PREDICTED_THREAD_COUNT);
// Create the pools
_multiplexerDataPool = new Pool <MultiplexerData>(PREDICTED_THREAD_COUNT, () => new MultiplexerData {
ManualResetEvent = new ManualResetEvent(false)
}, multiplexerData =>
{
multiplexerData.Message = null;
multiplexerData.ManualResetEvent.Reset();
});
_socketAsyncEventArgsSendPool = new Pool <SocketAsyncEventArgs>(PREDICTED_THREAD_COUNT, () =>
{
var poolItem = new SocketAsyncEventArgs();
poolItem.Completed += OperationCallback;
return(poolItem);
});
_socketAsyncEventArgsReceivePool = new Pool <SocketAsyncEventArgs>(PREDICTED_THREAD_COUNT, () =>
{
var poolItem = new SocketAsyncEventArgs();
poolItem.SetBuffer(new byte[messageBufferSize], 0, messageBufferSize);
poolItem.Completed += OperationCallback;
return(poolItem);
});
_socketAsyncEventArgsKeepAlivePool = new Pool <SocketAsyncEventArgs>(PREDICTED_THREAD_COUNT, () =>
{
var poolItem = new SocketAsyncEventArgs();
poolItem.SetBuffer(ProtocolHelper.ControlBytesPlaceholder, 0, ProtocolHelper.ControlBytesPlaceholder.Length);
poolItem.Completed += OperationCallback;
return(poolItem);
});
_receivedMessagePool = new Pool <ReceivedMessage>(PREDICTED_THREAD_COUNT, () => new ReceivedMessage(), receivedMessage =>
{
receivedMessage.Message = null;
receivedMessage.Socket = null;
});
_messageReceivedArgsPool = new Pool <MessageReceivedArgs>(PREDICTED_THREAD_COUNT, () => new MessageReceivedArgs(), messageReceivedArgs => { messageReceivedArgs.ReceivedMessage = null; });
_socketErrorArgsPool = new Pool <SocketErrorArgs>(PREDICTED_THREAD_COUNT, () => new SocketErrorArgs(), socketErrorArgs => { socketErrorArgs.Exception = null; });
}
private void ProcessReceivedMessage(ConnectedClient connectedClient)
{
int bytesToRead = -1;
int threadId = -1;
int availableTest = 0;
int controlBytesOffset = 0;
byte[] protocolBuffer = new byte[_controlBytesPlaceholder.Length];
byte[] resultBuffer = null;
var handler = connectedClient.Socket;
BlockingQueue <SocketAsyncEventArgs> receiveBufferQueue = null;
_currentlyConnectedClientsReceiveQueuesLock.EnterReadLock();
try
{
if (!_currentlyConnectedClientsReceiveQueues.TryGetValue(handler, out receiveBufferQueue))
{
// Peace out!
return;
}
}
finally
{
_currentlyConnectedClientsReceiveQueuesLock.ExitReadLock();
}
// Loop until socket is done
while (_isListening)
{
// If the socket is disposed, we're done
try
{
availableTest = handler.Available;
}
catch (ObjectDisposedException)
{
// Peace out!
return;
}
// Get the next buffer from the queue
var socketAsyncEventArgs = receiveBufferQueue.Dequeue();
if (socketAsyncEventArgs == null)
{
continue;
}
var buffer = socketAsyncEventArgs.Buffer;
int bytesRead = socketAsyncEventArgs.BytesTransferred;
int currentOffset = 0;
while (currentOffset < bytesRead)
{
// Check if we need to get our control byte values
if (bytesToRead == -1)
{
var controlBytesNeeded = _controlBytesPlaceholder.Length - controlBytesOffset;
var controlBytesAvailable = bytesRead - currentOffset;
var controlBytesToCopy = Math.Min(controlBytesNeeded, controlBytesAvailable);
// Copy bytes to control buffer
Buffer.BlockCopy(buffer, currentOffset, protocolBuffer, controlBytesOffset, controlBytesToCopy);
controlBytesOffset += controlBytesToCopy;
currentOffset += controlBytesToCopy;
// Check if done
if (controlBytesOffset == _controlBytesPlaceholder.Length)
{
// Parse out control bytes
ExtractControlBytes(protocolBuffer, out bytesToRead, out threadId);
// Reset control bytes offset
controlBytesOffset = 0;
// Ensure message is not larger than maximum message size
if (bytesToRead > _maxMessageSize)
{
HandleCommunicationError(handler, new InvalidOperationException(string.Format("message of length {0} exceeds maximum message length of {1}", bytesToRead, _maxMessageSize)));
return;
}
}
// Continue the loop
continue;
}
// Have control bytes, get message bytes
// SPECIAL CASE: if empty message, skip a bunch of stuff
if (bytesToRead != 0)
{
// Initialize buffer if needed
if (resultBuffer == null)
{
resultBuffer = new byte[bytesToRead];
}
var bytesAvailable = bytesRead - currentOffset;
var bytesToCopy = Math.Min(bytesToRead, bytesAvailable);
// Copy bytes to buffer
Buffer.BlockCopy(buffer, currentOffset, resultBuffer, resultBuffer.Length - bytesToRead, bytesToCopy);
currentOffset += bytesToCopy;
bytesToRead -= bytesToCopy;
}
// Check if we're done
if (bytesToRead == 0)
{
if (resultBuffer != null)
{
// Done, add to complete received messages
CompleteMessage(handler, threadId, resultBuffer);
// Reset message state
resultBuffer = null;
}
bytesToRead = -1;
threadId = -1;
connectedClient.LastResponse = DateTime.UtcNow;
}
}
// Push the buffer back onto the pool
_socketAsyncEventArgsReceivePool.Push(socketAsyncEventArgs);
}
}
private void ProcessReceivedMessage(MessageState messageState)
{
int currentOffset = 0;
int bytesRead = 0;
while (_isDoingSomething)
{
// Check if we need a buffer
if (messageState.Buffer == null)
{
// Get the next buffer
BlockingQueue <KeyValuePair <byte[], int> > queue = null;
_receiveBufferQueueLock.EnterReadLock();
try
{
if (!_receiveBufferQueue.TryGetValue(messageState.Handler.GetHashCode(), out queue))
{
throw new Exception("FATAL: No receive queue created for current socket");
}
}
finally
{
_receiveBufferQueueLock.ExitReadLock();
}
var receiveBufferEntry = queue.Dequeue();
messageState.Buffer = receiveBufferEntry.Key;
currentOffset = 0;
bytesRead = receiveBufferEntry.Value;
}
// Check if we need to get our control byte values
if (messageState.TotalBytesToRead == -1)
{
// We do, see if we have enough bytes received to get them
if (currentOffset + _controlBytesPlaceholder.Length > currentOffset + bytesRead)
{
// We don't yet have enough bytes to read the control bytes, so get more bytes
// Loop until we have enough data to proceed
int bytesNeeded = _controlBytesPlaceholder.Length - bytesRead;
while (bytesNeeded > 0)
{
// Combine the buffers
BlockingQueue <KeyValuePair <byte[], int> > queue = null;
_receiveBufferQueueLock.EnterReadLock();
try
{
if (!_receiveBufferQueue.TryGetValue(messageState.Handler.GetHashCode(), out queue))
{
throw new Exception("FATAL: No receive queue created for current socket");
}
}
finally
{
_receiveBufferQueueLock.ExitReadLock();
}
var nextBufferEntry = queue.Dequeue();
var combinedBuffer = new byte[bytesRead + nextBufferEntry.Value];
Buffer.BlockCopy(messageState.Buffer, currentOffset, combinedBuffer, 0, bytesRead);
Buffer.BlockCopy(nextBufferEntry.Key, 0, combinedBuffer, bytesRead, nextBufferEntry.Value);
// Set the new combined buffer and appropriate bytes read
messageState.Buffer = combinedBuffer;
// Reset bytes read and current offset
currentOffset = 0;
bytesRead = combinedBuffer.Length;
// Subtract from bytes needed
bytesNeeded -= nextBufferEntry.Value;
}
}
// Parse out control bytes
ExtractControlBytes(messageState.Buffer, currentOffset, out messageState.TotalBytesToRead, out messageState.ThreadId);
// Offset the index by the control bytes
currentOffset += _controlBytesPlaceholder.Length;
// Take control bytes off of bytes read
bytesRead -= _controlBytesPlaceholder.Length;
}
int numberOfBytesToRead = Math.Min(bytesRead, messageState.TotalBytesToRead);
messageState.Data.Write(messageState.Buffer, currentOffset, numberOfBytesToRead);
// Set total bytes read
int originalTotalBytesToRead = messageState.TotalBytesToRead;
messageState.TotalBytesToRead -= numberOfBytesToRead;
// Check if we're done
if (messageState.TotalBytesToRead == 0)
{
// Done, add to complete received messages
CompleteMessage(messageState.Handler, messageState.ThreadId, messageState.Data.ToArray());
}
// Check if we have an overlapping message frame in our message AKA if the bytesRead was larger than the total bytes to read
if (bytesRead > originalTotalBytesToRead)
{
// Get the number of bytes remaining to be read
int bytesRemaining = bytesRead - numberOfBytesToRead;
// Set total bytes to read to default
messageState.TotalBytesToRead = -1;
// Dispose and reinitialize data stream
messageState.Data.Dispose();
messageState.Data = new MemoryStream();
// Now we have the next message, so recursively process it
currentOffset += numberOfBytesToRead;
bytesRead = bytesRemaining;
continue;
}
// Only create a new message state if we are done with this message
if (!(bytesRead < originalTotalBytesToRead))
{
// Get new state for the next message but transfer over handler
Socket handler = messageState.Handler;
messageState.Data.Dispose();
messageState.Data = new MemoryStream();
messageState.Handler = handler;
messageState.TotalBytesToRead = -1;
messageState.ThreadId = -1;
}
// Reset buffer for next message
_bufferPool.Push(messageState.Buffer);
messageState.Buffer = null;
}
}