From 45ba8dff217d1fdb95dcbb46758319025503312b Mon Sep 17 00:00:00 2001
From: Mauro Servienti <mauro.servienti@gmail.com>
Date: Thu, 4 Jun 2020 19:48:54 +0200
Subject: [PATCH 1/5] Add current message to the timeout reschedule rule

---
 .../RescheduleTimeoutsBehavior.cs                     | 11 +++++++----
 1 file changed, 7 insertions(+), 4 deletions(-)

diff --git a/src/NServiceBus.IntegrationTesting/RescheduleTimeoutsBehavior.cs b/src/NServiceBus.IntegrationTesting/RescheduleTimeoutsBehavior.cs
index 06cba21e..6f60ec6d 100644
--- a/src/NServiceBus.IntegrationTesting/RescheduleTimeoutsBehavior.cs
+++ b/src/NServiceBus.IntegrationTesting/RescheduleTimeoutsBehavior.cs
@@ -19,14 +19,17 @@ public RescheduleTimeoutsBehavior(IntegrationScenarioContext integrationContext)
 
         public override Task Invoke(IOutgoingSendContext context, Func<Task> next)
         {
-            if (integrationContext.TryGetTimeoutRescheduleRule(context.Message.MessageType, out Func<DoNotDeliverBefore, DoNotDeliverBefore> rule))
+            if (integrationContext.TryGetTimeoutRescheduleRule(context.Message.MessageType, out Func<object, DoNotDeliverBefore, DoNotDeliverBefore> rule))
             {
                 var constraints = context.Extensions.Get<List<DeliveryConstraint>>();
                 var doNotDeliverBefore = constraints.OfType<DoNotDeliverBefore>().SingleOrDefault();
 
-                var newDoNotDeliverBefore = rule(doNotDeliverBefore);
-                constraints.Remove(doNotDeliverBefore);
-                constraints.Add(newDoNotDeliverBefore);
+                var newDoNotDeliverBefore = rule(context.Message, doNotDeliverBefore);
+                if(newDoNotDeliverBefore != doNotDeliverBefore)
+                {
+                   constraints.Remove(doNotDeliverBefore);
+                   constraints.Add(newDoNotDeliverBefore);
+                }
             }
 
             return next();

From 453fc1cbe56732fa46f3cd1e45e2970fa556f13a Mon Sep 17 00:00:00 2001
From: Mauro Servienti <mauro.servienti@gmail.com>
Date: Thu, 4 Jun 2020 19:50:36 +0200
Subject: [PATCH 2/5] Add message to the rule

---
 .../IntegrationScenarioContext.cs                           | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/src/NServiceBus.IntegrationTesting/IntegrationScenarioContext.cs b/src/NServiceBus.IntegrationTesting/IntegrationScenarioContext.cs
index 4a9f6766..7e7b2f77 100644
--- a/src/NServiceBus.IntegrationTesting/IntegrationScenarioContext.cs
+++ b/src/NServiceBus.IntegrationTesting/IntegrationScenarioContext.cs
@@ -13,7 +13,7 @@ public class IntegrationScenarioContext : ScenarioContext
         readonly ConcurrentBag<HandlerInvocation> invokedHandlers = new ConcurrentBag<HandlerInvocation>();
         readonly ConcurrentBag<SagaInvocation> invokedSagas = new ConcurrentBag<SagaInvocation>();
         readonly ConcurrentBag<OutgoingMessageOperation> outgoingMessageOperations = new ConcurrentBag<OutgoingMessageOperation>();
-        readonly Dictionary<Type, Func<DoNotDeliverBefore, DoNotDeliverBefore>> timeoutRescheduleRules = new Dictionary<Type, Func<DoNotDeliverBefore, DoNotDeliverBefore>>();
+        readonly Dictionary<Type, Func<object, DoNotDeliverBefore, DoNotDeliverBefore>> timeoutRescheduleRules = new Dictionary<Type, Func<object, DoNotDeliverBefore, DoNotDeliverBefore>>();
 
         public IEnumerable<HandlerInvocation> InvokedHandlers { get { return invokedHandlers; } }
         public IEnumerable<SagaInvocation> InvokedSagas { get { return invokedSagas; } }
@@ -26,7 +26,7 @@ internal HandlerInvocation CaptureInvokedHandler(HandlerInvocation invocation)
             return invocation;
         }
 
-        public void RegisterTimeoutRescheduleRule<TTimeout>(Func<DoNotDeliverBefore, DoNotDeliverBefore> rule)
+        public void RegisterTimeoutRescheduleRule<TTimeout>(Func<object, DoNotDeliverBefore, DoNotDeliverBefore> rule)
         {
             if (timeoutRescheduleRules.ContainsKey(typeof(TTimeout)))
             {
@@ -36,7 +36,7 @@ public void RegisterTimeoutRescheduleRule<TTimeout>(Func<DoNotDeliverBefore, DoN
             timeoutRescheduleRules.Add(typeof(TTimeout), rule);
         }
 
-        internal bool TryGetTimeoutRescheduleRule(Type timeoutMessageType, out Func<DoNotDeliverBefore, DoNotDeliverBefore> rule)
+        internal bool TryGetTimeoutRescheduleRule(Type timeoutMessageType, out Func<object, DoNotDeliverBefore, DoNotDeliverBefore> rule)
         {
             return timeoutRescheduleRules.TryGetValue(timeoutMessageType, out rule);
         }

From 28ddb70a0cb4a9d66135a20d4765a3305cb104b4 Mon Sep 17 00:00:00 2001
From: Mauro Servienti <mauro.servienti@gmail.com>
Date: Thu, 4 Jun 2020 19:50:48 +0200
Subject: [PATCH 3/5] Add message to the rule


From 8fabc7aee15554ce46e6f879b9160c23bf72d3cc Mon Sep 17 00:00:00 2001
From: Mauro Servienti <mauro.servienti@gmail.com>
Date: Thu, 4 Jun 2020 19:51:39 +0200
Subject: [PATCH 4/5] Add message to the rule

---
 src/MySystem.AcceptanceTests/When_requesting_a_timeout.cs | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/MySystem.AcceptanceTests/When_requesting_a_timeout.cs b/src/MySystem.AcceptanceTests/When_requesting_a_timeout.cs
index b851914b..2135d512 100644
--- a/src/MySystem.AcceptanceTests/When_requesting_a_timeout.cs
+++ b/src/MySystem.AcceptanceTests/When_requesting_a_timeout.cs
@@ -29,7 +29,7 @@ public async Task It_should_be_rescheduled_and_handled()
         {
             var context = await Scenario.Define<IntegrationScenarioContext>(ctx =>
             {
-                ctx.RegisterTimeoutRescheduleRule<ASaga.MyTimeout>(currentDelay =>
+                ctx.RegisterTimeoutRescheduleRule<ASaga.MyTimeout>((message, currentDelay) =>
                 {
                     return new DoNotDeliverBefore(DateTime.UtcNow.AddSeconds(5));
                 });

From 7d3e149a73b7d2280a164ab3fc3b9724d33e95c4 Mon Sep 17 00:00:00 2001
From: Mauro Servienti <mauro.servienti@gmail.com>
Date: Thu, 4 Jun 2020 20:39:32 +0200
Subject: [PATCH 5/5] Ooops

---
 .../Reschedule_Timeout_Behavior.cs                              | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/NServiceBus.IntegrationTesting.Tests/Reschedule_Timeout_Behavior.cs b/src/NServiceBus.IntegrationTesting.Tests/Reschedule_Timeout_Behavior.cs
index fcac744d..ef346148 100644
--- a/src/NServiceBus.IntegrationTesting.Tests/Reschedule_Timeout_Behavior.cs
+++ b/src/NServiceBus.IntegrationTesting.Tests/Reschedule_Timeout_Behavior.cs
@@ -19,7 +19,7 @@ public async Task Should_Reschedule_Timeouts()
             var expectedDeliveryAt = new DateTime(2020, 1, 1);
 
             var scenarioContext = new IntegrationScenarioContext();
-            scenarioContext.RegisterTimeoutRescheduleRule<AMessage>(currentDelay =>
+            scenarioContext.RegisterTimeoutRescheduleRule<AMessage>((msg, currentDelay) =>
             {
                 return new DoNotDeliverBefore(expectedDeliveryAt);
             });