private async Task<bool> Invoke(MessageResult result, Action<Subscription, object> beforeInvoke, object state)
{
// Change the state from idle to invoking callback
var prevState = Interlocked.CompareExchange(ref _subscriptionState,
SubscriptionState.InvokingCallback,
SubscriptionState.Idle);
if (prevState == SubscriptionState.Disposed)
{
// Only allow terminal messages after dispose
if (!result.Terminal)
{
return false;
}
}
beforeInvoke(this, state);
_counters.MessageBusMessagesReceivedTotal.IncrementBy(result.TotalCount);
_counters.MessageBusMessagesReceivedPerSec.IncrementBy(result.TotalCount);
try
{
return await _callback(result, _callbackState);
}
finally
{
// Go from invoking callback to idle
Interlocked.CompareExchange(ref _subscriptionState,
SubscriptionState.Idle,
SubscriptionState.InvokingCallback);
}
}
private Task<bool> Invoke(MessageResult result, Action<object> beforeInvoke, object state)
{
// Change the state from idle to invoking callback
var prevState = Interlocked.CompareExchange(ref _subscriptionState,
SubscriptionState.InvokingCallback,
SubscriptionState.Idle);
if (prevState == SubscriptionState.Disposed)
{
// Only allow terminal messages after dispose
if (!result.Terminal)
{
return TaskAsyncHelper.False;
}
}
beforeInvoke(state);
_counters.MessageBusMessagesReceivedTotal.IncrementBy(result.TotalCount);
_counters.MessageBusMessagesReceivedPerSec.IncrementBy(result.TotalCount);
return _callback.Invoke(result, _callbackState).ContinueWith(task =>
{
// Go from invoking callback to idle
Interlocked.CompareExchange(ref _subscriptionState,
SubscriptionState.Idle,
SubscriptionState.InvokingCallback);
return task;
},
TaskContinuationOptions.ExecuteSynchronously).FastUnwrap();
}
private Task<bool> TriggerAcks(MessageResult result, object state)
{
result.Messages.Enumerate<object>(m => m.IsAck,
(s, m) => ((IAckHandler)s).TriggerAck(m.CommandId),
state: _ackHandler);
return TaskAsyncHelper.True;
}
private void WorkImpl(TaskCompletionSource<object> taskCompletionSource)
{
Process:
if (!Alive)
{
// If this subscription is dead then return immediately
taskCompletionSource.TrySetResult(null);
return;
}
var items = new List<ArraySegment<Message>>();
int totalCount;
object state;
PerformWork(items, out totalCount, out state);
if (items.Count > 0)
{
var messageResult = new MessageResult(items, totalCount);
Task<bool> callbackTask = Invoke(messageResult, s => BeforeInvoke(s), state);
if (callbackTask.IsCompleted)
{
try
{
// Make sure exceptions propagate
callbackTask.Wait();
if (callbackTask.Result)
{
// Sync path
goto Process;
}
else
{
// If we're done pumping messages through to this subscription
// then dispose
Dispose();
// If the callback said it's done then stop
taskCompletionSource.TrySetResult(null);
}
}
catch (Exception ex)
{
if (ex.InnerException is TaskCanceledException)
{
taskCompletionSource.TrySetCanceled();
}
else
{
taskCompletionSource.TrySetUnwrappedException(ex);
}
}
}
else
{
WorkImplAsync(callbackTask, taskCompletionSource);
}
}
else
{
taskCompletionSource.TrySetResult(null);
}
}
public Task<bool> InvokeCallback(MessageResult result)
{
var response = _connection.GetResponse(result);
return _callback(response, _callbackState);
}
private PersistentResponse GetResponse(MessageResult result)
{
// Do a single sweep through the results to process commands and extract values
ProcessResults(result);
Debug.Assert(WriteCursor != null, "Unable to resolve the cursor since the method is null");
var response = new PersistentResponse(ExcludeMessage, WriteCursor);
response.Terminal = result.Terminal;
if (!result.Terminal)
{
// Only set these properties if the message isn't terminal
response.Messages = result.Messages;
response.Disconnect = _disconnected;
response.Aborted = _aborted;
response.TotalCount = result.TotalCount;
}
PopulateResponseState(response);
_counters.ConnectionMessagesReceivedTotal.IncrementBy(result.TotalCount);
_counters.ConnectionMessagesReceivedPerSec.IncrementBy(result.TotalCount);
return response;
}
private void ProcessResults(MessageResult result)
{
result.Messages.Enumerate<object>(message => message.IsAck || message.IsCommand,
(state, message) =>
{
if (message.IsAck)
{
_ackHandler.TriggerAck(message.CommandId);
}
else if (message.IsCommand)
{
var command = _serializer.Parse<Command>(message.Value, message.Encoding);
if (command == null)
{
if (MonoUtility.IsRunningMono)
{
return;
}
throw new SerializationException("Couldn't parse message " + message.Value);
}
ProcessCommand(command);
// Only send the ack if this command is waiting for it
if (message.WaitForAck)
{
// If we're on the same box and there's a pending ack for this command then
// just trip it
if (!_ackHandler.TriggerAck(message.CommandId))
{
_bus.Ack(_connectionId, message.CommandId).Catch();
}
}
}
}, null);
}
public virtual Task<bool> Invoke(MessageResult result)
{
return Invoke(result, () => { });
}
private static Task<bool> MessageBusCallback(MessageResult result, object state)
{
var context = (ReceiveContext)state;
return context.InvokeCallback(result);
}
public async Task Work()
{
// Set the state to working
Interlocked.Exchange(ref _state, State.Working);
while (Alive)
{
var items = new List<ArraySegment<Message>>();
int totalCount;
object state;
PerformWork(items, out totalCount, out state);
if (items.Count > 0)
{
var messageResult = new MessageResult(items, totalCount);
bool result = await Invoke(messageResult, (s, o) => s.BeforeInvoke(o), state);
if (!result)
{
Dispose();
// If the callback said it's done then stop
break;
}
}
else
{
break;
}
}
}
private void WorkImpl(TaskCompletionSource <object> taskCompletionSource)
{
Process:
if (!Alive)
{
// If this subscription is dead then return immediately
taskCompletionSource.TrySetResult(null);
return;
}
var items = new List <ArraySegment <Message> >();
int totalCount;
object state;
PerformWork(items, out totalCount, out state);
if (items.Count > 0)
{
var messageResult = new MessageResult(items, totalCount);
Task <bool> callbackTask = Invoke(messageResult, s => BeforeInvoke(s), state);
if (callbackTask.IsCompleted)
{
try
{
// Make sure exceptions propagate
callbackTask.Wait();
if (callbackTask.Result)
{
// Sync path
goto Process;
}
else
{
// If we're done pumping messages through to this subscription
// then dispose
Dispose();
// If the callback said it's done then stop
taskCompletionSource.TrySetResult(null);
}
}
catch (Exception ex)
{
if (ex.InnerException is TaskCanceledException)
{
taskCompletionSource.TrySetCanceled();
}
else
{
taskCompletionSource.TrySetUnwrappedException(ex);
}
}
}
else
{
WorkImplAsync(callbackTask, taskCompletionSource);
}
}
else
{
taskCompletionSource.TrySetResult(null);
}
}
private PersistentResponse GetResponse(MessageResult result)
{
// Do a single sweep through the results to process commands and extract values
ProcessResults(result);
Debug.Assert(GetCursor != null, "Unable to resolve the cursor since the method is null");
// Resolve the cursor
string id = GetCursor();
var response = new PersistentResponse(ExcludeMessage)
{
MessageId = id,
Messages = result.Messages,
Disconnect = _disconnected,
Aborted = _aborted,
TotalCount = result.TotalCount,
};
PopulateResponseState(response);
_counters.ConnectionMessagesReceivedTotal.IncrementBy(result.TotalCount);
_counters.ConnectionMessagesReceivedPerSec.IncrementBy(result.TotalCount);
return response;
}
private void ProcessResults(MessageResult result)
{
result.Messages.Enumerate(message => message.IsAck || message.IsCommand,
message =>
{
if (message.IsAck)
{
_ackHandler.TriggerAck(message.CommandId);
}
else if (message.IsCommand)
{
var command = _serializer.Parse<Command>(message.Value);
ProcessCommand(command);
// Only send the ack if this command is waiting for it
if (message.WaitForAck)
{
// If we're on the same box and there's a pending ack for this command then
// just trip it
if (!_ackHandler.TriggerAck(message.CommandId))
{
_bus.Ack(_connectionId, message.CommandId).Catch();
}
}
}
});
}
private Task<bool> HandleServerCommands(MessageResult result, object state)
{
result.Messages.Enumerate<object>(m => ServerSignal.Equals(m.Key),
(s, m) =>
{
var command = _serializer.Parse<ServerCommand>(m.Value, m.Encoding);
OnCommand(command);
},
state: null);
return TaskAsyncHelper.True;
}
private void ProcessResults(MessageResult result)
{
result.Messages.Enumerate<Connection>(message => message.IsCommand,
(connection, message) => ProcessResultsCore(connection, message), this);
}
private Task<bool> Invoke(MessageResult result, Action beforeInvoke)
{
// Change the state from idle to invoking callback
var state = Interlocked.CompareExchange(ref _subscriptionState,
SubscriptionState.InvokingCallback,
SubscriptionState.Idle);
if (state == SubscriptionState.Disposed)
{
// Only allow terminal messages after dispose
if (!result.Terminal)
{
return TaskAsyncHelper.False;
}
}
beforeInvoke();
return _callback.Invoke(result).ContinueWith(task =>
{
// Go from invoking callback to idle
Interlocked.CompareExchange(ref _subscriptionState,
SubscriptionState.Idle,
SubscriptionState.InvokingCallback);
if (task.IsFaulted)
{
return TaskAsyncHelper.FromError<bool>(task.Exception);
}
return TaskAsyncHelper.FromResult(task.Result);
},
TaskContinuationOptions.ExecuteSynchronously).FastUnwrap();
}
public virtual Task<bool> Invoke(MessageResult result)
{
return Invoke(result, state => { }, state: null);
}
public void AcksAreSentToAckSubscriber()
{
// Arrange
var waitCommand = new Command { WaitForAck = true };
var ackerId = "acker";
var waiterId = "waiter";
var messageId = "messageId";
var maxMessages = 1;
var ackHandler = new Mock<IAckHandler>();
ackHandler.Setup(m => m.TriggerAck(waitCommand.Id)).Returns(false);
var messageBus = new Mock<IMessageBus>();
Message waitMessage = null;
Message ackMessage = null;
messageBus.Setup(m => m.Publish(It.IsAny<Message>())).Returns<Message>(m =>
{
if (m.WaitForAck)
{
waitMessage = m;
}
else if (m.IsAck)
{
ackMessage = m;
}
return TaskAsyncHelper.Empty;
});
var loggerFactory = new Mock<ILoggerFactory>();
var counters = new PerformanceCounterManager(loggerFactory.Object);
var serializer = JsonUtility.CreateDefaultSerializer();
var waiterConnection = new Connection(messageBus.Object,
serializer,
"signal",
waiterId,
new string[] { },
new string[] { },
loggerFactory.Object,
ackHandler.Object,
counters,
new Mock<IProtectedData>().Object,
new MemoryPool());
// Act
waiterConnection.Send(ackerId, waitCommand);
// Assert
Assert.NotNull(waitMessage);
Assert.Equal(waiterId, waitMessage.Source);
Assert.Equal(PrefixHelper.GetConnectionId(ackerId), waitMessage.Key);
// Arrange some more now that we have a waitMessage
var messages = new List<ArraySegment<Message>>()
{
new ArraySegment<Message>(new[] { waitMessage })
};
var messageResult = new MessageResult(messages, 1);
var ackerConnection = new Connection(messageBus.Object,
serializer,
"signal",
ackerId,
new string[] { },
new string[] { },
loggerFactory.Object,
ackHandler.Object,
counters,
new Mock<IProtectedData>().Object,
new Mock<IMemoryPool>().Object);
ackerConnection.WriteCursor = _ => { };
messageBus.Setup(m => m.Subscribe(ackerConnection,
messageId,
It.IsAny<Func<MessageResult, object, Task<bool>>>(),
maxMessages,
It.IsAny<object>()))
.Callback<ISubscriber,
string,
Func<MessageResult, object, Task<bool>>,
int,
object>((subsciber, cursor, callback, max, state) =>
{
callback(messageResult, state);
});
// Act
ackerConnection.Receive(messageId, (_, __) => TaskAsyncHelper.False, maxMessages, null);
// Assert
Assert.NotNull(ackMessage);
Assert.Equal(ackerId, ackMessage.Source);
Assert.Equal(AckSubscriber.Signal, ackMessage.Key);
}
private Task<bool> HandleServerCommands(MessageResult result)
{
result.Messages.Enumerate(m => ServerSignal.Equals(m.Key),
m =>
{
var command = _serializer.Parse<ServerCommand>(m.Value);
OnCommand(command);
});
return TaskAsyncHelper.True;
}
public virtual Task <bool> Invoke(MessageResult result)
{
return(Invoke(result, (s, o) => { }, state: null));
}
