Skip to content

Commit df40d1c

Browse files
authored
Support context propagators in worker factory options
1 parent 0dd64a7 commit df40d1c

5 files changed

Lines changed: 145 additions & 1 deletion

File tree

‎temporal-sdk/src/main/java/io/temporal/common/SimplePlugin.java‎

Lines changed: 9 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -285,6 +285,15 @@ public void configureWorkerFactory(@Nonnull WorkerFactoryOptions.Builder builder
285285
combined.addAll(workerInterceptors);
286286
builder.setWorkerInterceptors(combined.toArray(new WorkerInterceptor[0]));
287287
}
288+
289+
// Add context propagators
290+
if (!contextPropagators.isEmpty()) {
291+
List<ContextPropagator> existing = builder.build().getContextPropagators();
292+
List<ContextPropagator> combined =
293+
new ArrayList<>(existing != null ? existing : new ArrayList<>());
294+
combined.addAll(contextPropagators);
295+
builder.setContextPropagators(combined);
296+
}
288297
}
289298

290299
@Override

‎temporal-sdk/src/main/java/io/temporal/worker/WorkerFactory.java‎

Lines changed: 13 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -10,6 +10,7 @@
1010
import io.temporal.api.workflowservice.v1.DescribeNamespaceResponse;
1111
import io.temporal.client.WorkflowClient;
1212
import io.temporal.client.WorkflowClientOptions;
13+
import io.temporal.common.context.ContextPropagator;
1314
import io.temporal.common.converter.DataConverter;
1415
import io.temporal.internal.client.WorkflowClientInternal;
1516
import io.temporal.internal.common.PluginUtils;
@@ -27,6 +28,7 @@
2728
import java.util.Arrays;
2829
import java.util.Collections;
2930
import java.util.HashMap;
31+
import java.util.LinkedHashMap;
3032
import java.util.List;
3133
import java.util.Map;
3234
import java.util.Objects;
@@ -192,6 +194,16 @@ public synchronized Worker newWorker(String taskQueue, WorkerOptions options) {
192194
// Apply plugin configuration to worker options (forward order)
193195
options = applyWorkerPluginConfiguration(taskQueue, options, this.plugins);
194196

197+
Map<String, ContextPropagator> contextPropagatorsByName = new LinkedHashMap<>();
198+
for (ContextPropagator propagator : workflowClient.getOptions().getContextPropagators()) {
199+
contextPropagatorsByName.putIfAbsent(propagator.getName(), propagator);
200+
}
201+
for (ContextPropagator propagator : factoryOptions.getContextPropagators()) {
202+
contextPropagatorsByName.putIfAbsent(propagator.getName(), propagator);
203+
}
204+
List<ContextPropagator> contextPropagators =
205+
new ArrayList<>(contextPropagatorsByName.values());
206+
195207
// Only one worker can exist for a task queue
196208
Worker existingWorker = workers.get(taskQueue);
197209
if (existingWorker == null) {
@@ -206,7 +218,7 @@ public synchronized Worker newWorker(String taskQueue, WorkerOptions options) {
206218
cache,
207219
true,
208220
workflowThreadExecutor,
209-
workflowClient.getOptions().getContextPropagators(),
221+
contextPropagators,
210222
plugins,
211223
((WorkflowClientInternal) workflowClient.getInternal()).getWorkerGroupingKey(),
212224
namespaceCapabilities);

‎temporal-sdk/src/main/java/io/temporal/worker/WorkerFactoryOptions.java‎

Lines changed: 28 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -3,8 +3,11 @@
33
import com.google.common.annotations.VisibleForTesting;
44
import com.google.common.base.Preconditions;
55
import io.temporal.common.Experimental;
6+
import io.temporal.common.context.ContextPropagator;
67
import io.temporal.common.interceptors.WorkerInterceptor;
78
import java.time.Duration;
9+
import java.util.Collections;
10+
import java.util.List;
811
import java.util.concurrent.ExecutorService;
912
import javax.annotation.Nullable;
1013

@@ -38,6 +41,7 @@ public static class Builder {
3841
private int workflowCacheSize;
3942
private int maxWorkflowThreadCount;
4043
private WorkerInterceptor[] workerInterceptors;
44+
private List<ContextPropagator> contextPropagators;
4145
private WorkerPlugin[] plugins;
4246
private boolean enableLoggingInReplay;
4347
private boolean usingVirtualWorkflowThreads;
@@ -55,6 +59,7 @@ private Builder(WorkerFactoryOptions options) {
5559
this.workflowCacheSize = options.workflowCacheSize;
5660
this.maxWorkflowThreadCount = options.maxWorkflowThreadCount;
5761
this.workerInterceptors = options.workerInterceptors;
62+
this.contextPropagators = options.contextPropagators;
5863
this.plugins = options.plugins;
5964
this.enableLoggingInReplay = options.enableLoggingInReplay;
6065
this.usingVirtualWorkflowThreads = options.usingVirtualWorkflowThreads;
@@ -106,6 +111,15 @@ public Builder setWorkerInterceptors(WorkerInterceptor... workerInterceptors) {
106111
return this;
107112
}
108113

114+
/**
115+
* Sets the context propagators to use with workers created by this factory. These are appended
116+
* to the propagators configured on the workflow client.
117+
*/
118+
public Builder setContextPropagators(List<ContextPropagator> contextPropagators) {
119+
this.contextPropagators = contextPropagators;
120+
return this;
121+
}
122+
109123
/**
110124
* Sets the worker plugins to use with workers created by this factory. Plugins can modify
111125
* worker configuration and wrap worker lifecycle.
@@ -179,6 +193,7 @@ public WorkerFactoryOptions build() {
179193
maxWorkflowThreadCount,
180194
workflowHostLocalTaskQueueScheduleToStartTimeout,
181195
workerInterceptors,
196+
contextPropagators,
182197
plugins,
183198
enableLoggingInReplay,
184199
usingVirtualWorkflowThreads,
@@ -204,6 +219,7 @@ public WorkerFactoryOptions validateAndBuildWithDefaults() {
204219
maxWorkflowThreadCount,
205220
workflowHostLocalTaskQueueScheduleToStartTimeout,
206221
workerInterceptors == null ? new WorkerInterceptor[0] : workerInterceptors,
222+
contextPropagators == null ? Collections.emptyList() : contextPropagators,
207223
plugins == null ? new WorkerPlugin[0] : plugins,
208224
enableLoggingInReplay,
209225
usingVirtualWorkflowThreads,
@@ -217,6 +233,7 @@ public WorkerFactoryOptions validateAndBuildWithDefaults() {
217233
private final int maxWorkflowThreadCount;
218234
private final @Nullable Duration workflowHostLocalTaskQueueScheduleToStartTimeout;
219235
private final WorkerInterceptor[] workerInterceptors;
236+
private final List<ContextPropagator> contextPropagators;
220237
private final WorkerPlugin[] plugins;
221238
private final boolean enableLoggingInReplay;
222239
private final boolean usingVirtualWorkflowThreads;
@@ -228,6 +245,7 @@ private WorkerFactoryOptions(
228245
int maxWorkflowThreadCount,
229246
@Nullable Duration workflowHostLocalTaskQueueScheduleToStartTimeout,
230247
WorkerInterceptor[] workerInterceptors,
248+
List<ContextPropagator> contextPropagators,
231249
WorkerPlugin[] plugins,
232250
boolean enableLoggingInReplay,
233251
boolean usingVirtualWorkflowThreads,
@@ -268,6 +286,7 @@ private WorkerFactoryOptions(
268286
this.workflowHostLocalTaskQueueScheduleToStartTimeout =
269287
workflowHostLocalTaskQueueScheduleToStartTimeout;
270288
this.workerInterceptors = workerInterceptors;
289+
this.contextPropagators = contextPropagators;
271290
this.plugins = plugins;
272291
this.enableLoggingInReplay = enableLoggingInReplay;
273292
this.usingVirtualWorkflowThreads = usingVirtualWorkflowThreads;
@@ -292,6 +311,15 @@ public WorkerInterceptor[] getWorkerInterceptors() {
292311
return workerInterceptors;
293312
}
294313

314+
/**
315+
* Returns the context propagators used with workers created by this factory.
316+
*
317+
* @return the list of context propagators, never null
318+
*/
319+
public List<ContextPropagator> getContextPropagators() {
320+
return contextPropagators;
321+
}
322+
295323
/**
296324
* Returns the worker plugins configured for this factory.
297325
*

‎temporal-sdk/src/test/java/io/temporal/common/SimplePluginBuilderTest.java‎

Lines changed: 13 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -30,6 +30,7 @@
3030
import io.temporal.client.WorkflowClientOptions;
3131
import io.temporal.client.schedules.ScheduleClientOptions;
3232
import io.temporal.client.schedules.ScheduleClientPlugin;
33+
import io.temporal.common.context.ContextPropagator;
3334
import io.temporal.common.converter.DataConverter;
3435
import io.temporal.common.interceptors.ActivityClientInterceptor;
3536
import io.temporal.common.interceptors.ActivityClientInterceptorBase;
@@ -88,6 +89,18 @@ public void testAddWorkerInterceptors() {
8889
assertSame(interceptor, interceptors[0]);
8990
}
9091

92+
@Test
93+
public void testAddContextPropagatorsToWorkerFactory() {
94+
ContextPropagator propagator = mock(ContextPropagator.class);
95+
96+
SimplePlugin plugin = SimplePlugin.newBuilder("test").addContextPropagators(propagator).build();
97+
98+
WorkerFactoryOptions.Builder builder = WorkerFactoryOptions.newBuilder();
99+
((io.temporal.worker.WorkerPlugin) plugin).configureWorkerFactory(builder);
100+
101+
assertEquals(Collections.singletonList(propagator), builder.build().getContextPropagators());
102+
}
103+
91104
@Test
92105
public void testAddClientInterceptors() {
93106
WorkflowClientInterceptor interceptor = new WorkflowClientInterceptorBase() {};

‎temporal-sdk/src/test/java/io/temporal/workflow/ContextPropagationTest.java‎

Lines changed: 82 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,6 +14,7 @@
1414
import io.temporal.testing.TestEnvironmentOptions;
1515
import io.temporal.testing.TestWorkflowEnvironment;
1616
import io.temporal.worker.Worker;
17+
import io.temporal.worker.WorkerFactoryOptions;
1718
import io.temporal.workflow.shared.TestActivities;
1819
import io.temporal.workflow.shared.TestWorkflows;
1920
import java.time.Duration;
@@ -99,6 +100,77 @@ public void testThreadContextPropagation() {
99100
assertEquals("asynctesting123", result);
100101
}
101102

103+
@Test
104+
public void testThreadContextPropagationFromWorkerFactoryOptions() {
105+
TestEnvironmentOptions options =
106+
TestEnvironmentOptions.newBuilder()
107+
.setWorkerFactoryOptions(
108+
WorkerFactoryOptions.newBuilder()
109+
.setContextPropagators(Collections.singletonList(new TestContextPropagator()))
110+
.build())
111+
.build();
112+
TestWorkflowEnvironment environment = TestWorkflowEnvironment.newInstance(options);
113+
try {
114+
Worker worker = environment.newWorker(TASK_QUEUE);
115+
worker.registerWorkflowImplementationTypes(FactoryContextPropagationThreadWorkflowImpl.class);
116+
environment.start();
117+
118+
WorkflowOptions workflowOptions =
119+
WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build();
120+
TestWorkflows.TestWorkflow1 workflow =
121+
environment
122+
.getWorkflowClient()
123+
.newWorkflowStub(TestWorkflows.TestWorkflow1.class, workflowOptions);
124+
assertEquals("testing123", workflow.execute("testing123"));
125+
} finally {
126+
environment.close();
127+
}
128+
}
129+
130+
@Test
131+
public void testDuplicateContextPropagatorsAreIgnored() {
132+
TestContextPropagator clientPropagator = new TestContextPropagator();
133+
TestContextPropagator factoryPropagator =
134+
new TestContextPropagator() {
135+
@Override
136+
public String getName() {
137+
return clientPropagator.getName();
138+
}
139+
140+
@Override
141+
public Object getCurrentContext() {
142+
throw new AssertionError("Duplicate factory context propagator was used");
143+
}
144+
};
145+
TestEnvironmentOptions options =
146+
TestEnvironmentOptions.newBuilder()
147+
.setWorkflowClientOptions(
148+
WorkflowClientOptions.newBuilder()
149+
.setContextPropagators(Collections.singletonList(clientPropagator))
150+
.build())
151+
.setWorkerFactoryOptions(
152+
WorkerFactoryOptions.newBuilder()
153+
.setContextPropagators(Collections.singletonList(factoryPropagator))
154+
.build())
155+
.build();
156+
TestWorkflowEnvironment environment = TestWorkflowEnvironment.newInstance(options);
157+
try {
158+
Worker worker = environment.newWorker(TASK_QUEUE);
159+
worker.registerWorkflowImplementationTypes(FactoryContextPropagationThreadWorkflowImpl.class);
160+
environment.start();
161+
162+
WorkflowOptions workflowOptions =
163+
WorkflowOptions.newBuilder().setTaskQueue(TASK_QUEUE).build();
164+
TestWorkflows.TestWorkflow1 workflow =
165+
environment
166+
.getWorkflowClient()
167+
.newWorkflowStub(TestWorkflows.TestWorkflow1.class, workflowOptions);
168+
assertEquals("testing123", workflow.execute("testing123"));
169+
} finally {
170+
environment.close();
171+
}
172+
}
173+
102174
@Test
103175
public void testActivityContextPropagation() {
104176
Worker worker = testEnvironment.newWorker(TASK_QUEUE);
@@ -334,6 +406,16 @@ private String async() {
334406
}
335407
}
336408

409+
public static class FactoryContextPropagationThreadWorkflowImpl
410+
implements TestWorkflows.TestWorkflow1 {
411+
412+
@Override
413+
public String execute(String input) {
414+
MDC.put("test", input);
415+
return Async.function(() -> MDC.get("test")).get();
416+
}
417+
}
418+
337419
public static class ContextActivityImpl implements TestActivities.TestActivity1 {
338420
@Override
339421
public String execute(String input) {

0 commit comments

Comments
 (0)