Adding extra WithCreateSendMessage step in ConsumerDataflow.

houseofcat · Apr 30, 2024 · 36dd565 · 36dd565
1 parent 255669d
commit 36dd565
Show file tree

Hide file tree

Showing 3 changed files with 18 additions and 5 deletions.
diff --git a/guides/rabbitmq/ConsumerDataflow.md b/guides/rabbitmq/ConsumerDataflow.md
@@ -358,7 +358,7 @@ outbound message and assign it to the State.
 
 ```csharp
 dataflow.WithCreateSendMessage(
-    async (state) =>
+    (state) =>
     {
         var message = new Message
         {
@@ -718,7 +718,7 @@ var dataflowService = new ConsumerDataflowService<CustomWorkState>(rabbitService
 
 // Manually modify the internal Dataflow.
 dataflowService.Dataflow.WithCreateSendMessage(
-    async (state) =>
+    (state) =>
     {
         var message = new Message
         {

diff --git a/src/HouseofCat.RabbitMQ/Dataflows/ConsumerDataflow.cs b/src/HouseofCat.RabbitMQ/Dataflows/ConsumerDataflow.cs
@@ -355,6 +355,21 @@ public ConsumerDataflow<TState> WithCreateSendMessage(
         return this;
     }
 
+    public ConsumerDataflow<TState> WithCreateSendMessage(
+        Func<TState, TState> createMessage,
+        int? maxDoP = null,
+        bool? ensureOrdered = null,
+        int? boundedCapacity = null,
+        TaskScheduler taskScheduler = null)
+    {
+        if (_createSendMessage is null)
+        {
+            var executionOptions = GetExecuteStepOptions(maxDoP, ensureOrdered, boundedCapacity, taskScheduler ?? _taskScheduler);
+            _createSendMessage = GetWrappedTransformBlock(createMessage, executionOptions, GetSpanName("create send message"));
+        }
+        return this;
+    }
+
     public ConsumerDataflow<TState> WithSendCompressedStep(
         int? maxDoP = null,
         bool? ensureOrdered = null,

diff --git a/tests/RabbitMQ.ConsumerDataflowService/Program.cs b/tests/RabbitMQ.ConsumerDataflowService/Program.cs
@@ -28,7 +28,7 @@
 
 // Manually modify the internal Dataflow.
 dataflowService.Dataflow.WithCreateSendMessage(
-    async (state) =>
+    (state) =>
     {
         var message = new Message
         {
@@ -42,8 +42,6 @@
             ParentSpanContext = state.WorkflowSpan?.Context,
         };
 
-        await rabbitService.ComcryptAsync(message);
-
         state.SendMessage = message;
         return state;
     });