added ThreadFactoryWrapper in WorkerFactoryOptions

Author: shijiesheng

src/main/java/com/uber/cadence/internal/sync/SyncActivityWorker.java

@@ -31,10 +31,13 @@ public class SyncActivityWorker implements SuspendableWorker {
 
   private final ActivityWorker worker;
   private final POJOActivityTaskHandler taskHandler;
-  private final ScheduledExecutorService heartbeatExecutor = Executors.newScheduledThreadPool(4);
+  private final ScheduledExecutorService heartbeatExecutor;
 
   public SyncActivityWorker(
       IWorkflowService service, String domain, String taskList, SingleWorkerOptions options) {
+    heartbeatExecutor =
+        Executors.newScheduledThreadPool(
+            4, options.getThreadFactoryWrapper().wrap(Executors.privilegedThreadFactory()));
     taskHandler =
         new POJOActivityTaskHandler(service, domain, options.getDataConverter(), heartbeatExecutor);
     worker = new ActivityWorker(service, domain, taskList, options, taskHandler);

src/main/java/com/uber/cadence/internal/sync/SyncWorkflowWorker.java

@@ -51,8 +51,8 @@ public class SyncWorkflowWorker
   private final POJOWorkflowImplementationFactory factory;
   private final DataConverter dataConverter;
   private final POJOActivityTaskHandler laTaskHandler;
-  private final ScheduledExecutorService heartbeatExecutor = Executors.newScheduledThreadPool(4);
-  private final ScheduledExecutorService ldaHeartbeatExecutor = Executors.newScheduledThreadPool(4);
+  private final ScheduledExecutorService heartbeatExecutor;
+  private final ScheduledExecutorService ldaHeartbeatExecutor;
   private SuspendableWorker ldaWorker;
   private POJOActivityTaskHandler ldaTaskHandler;
   private final IWorkflowService service;
@@ -73,6 +73,16 @@ public SyncWorkflowWorker(
     this.dataConverter = workflowOptions.getDataConverter();
     this.service = service;
 
+    // heartbeat executors
+    heartbeatExecutor =
+        Executors.newScheduledThreadPool(
+            4,
+            localActivityOptions.getThreadFactoryWrapper().wrap(Executors.defaultThreadFactory()));
+    ldaHeartbeatExecutor =
+        Executors.newScheduledThreadPool(
+            4,
+            localActivityOptions.getThreadFactoryWrapper().wrap(Executors.defaultThreadFactory()));
+
     factory =
         new POJOWorkflowImplementationFactory(
             workflowOptions.getDataConverter(),

src/main/java/com/uber/cadence/internal/worker/ActivityWorker.java

@@ -100,7 +100,8 @@ public void start() {
               getOrCreateActivityPollTask(),
               new PollTaskExecutor<>(domain, taskList, options, new TaskHandlerImpl(handler)),
               options.getPollerOptions(),
-              options.getMetricsScope());
+              options.getMetricsScope(),
+              options.getThreadFactoryWrapper());
       poller.start();
       setPoller(poller);
       options.getMetricsScope().counter(MetricsType.WORKER_START_COUNTER).inc(1);

src/main/java/com/uber/cadence/internal/worker/LocalActivityWorker.java

@@ -83,7 +83,8 @@ public void start() {
               laPollTask,
               new PollTaskExecutor<>(domain, taskList, options, new TaskHandlerImpl(handler)),
               options.getPollerOptions(),
-              options.getMetricsScope());
+              options.getMetricsScope(),
+              options.getThreadFactoryWrapper());
       poller.start();
       setPoller(poller);
       options.getMetricsScope().counter(MetricsType.WORKER_START_COUNTER).inc(1);

src/main/java/com/uber/cadence/internal/worker/PollTaskExecutor.java

@@ -55,9 +55,15 @@ public interface TaskHandler<TT> {
             TimeUnit.SECONDS,
             new SynchronousQueue<>());
     taskExecutor.setThreadFactory(
-        new ExecutorThreadFactory(
-            options.getPollerOptions().getPollThreadNamePrefix().replaceFirst("Poller", "Executor"),
-            options.getPollerOptions().getUncaughtExceptionHandler()));
+        options
+            .getThreadFactoryWrapper()
+            .wrap(
+                new ExecutorThreadFactory(
+                    options
+                        .getPollerOptions()
+                        .getPollThreadNamePrefix()
+                        .replaceFirst("Poller", "Executor"),
+                    options.getPollerOptions().getUncaughtExceptionHandler())));
     taskExecutor.setRejectedExecutionHandler(new BlockCallerPolicy());
   }
 

src/main/java/com/uber/cadence/internal/worker/Poller.java

@@ -22,6 +22,7 @@
 import com.uber.cadence.internal.metrics.MetricsType;
 import com.uber.cadence.internal.worker.autoscaler.AutoScaler;
 import com.uber.cadence.internal.worker.autoscaler.AutoScalerFactory;
+import com.uber.cadence.worker.ThreadFactoryWrapper;
 import com.uber.m3.tally.Scope;
 import java.util.Objects;
 import java.util.concurrent.ArrayBlockingQueue;
@@ -72,24 +73,29 @@ interface ThrowingRunnable {
 
   private final AutoScaler pollerAutoScaler;
 
+  private final ThreadFactoryWrapper threadFactoryWrapper;
+
   public Poller(
       String identity,
       PollTask<T> pollTask,
       ShutdownableTaskExecutor<T> taskExecutor,
       PollerOptions pollerOptions,
-      Scope metricsScope) {
+      Scope metricsScope,
+      ThreadFactoryWrapper threadFactoryWrapper) {
     Objects.requireNonNull(identity, "identity cannot be null");
     Objects.requireNonNull(pollTask, "poll service should not be null");
     Objects.requireNonNull(taskExecutor, "taskExecutor should not be null");
     Objects.requireNonNull(pollerOptions, "pollerOptions should not be null");
     Objects.requireNonNull(metricsScope, "metricsScope should not be null");
+    Objects.requireNonNull(metricsScope, "threadFactoryWrapper should not be null");
 
     this.identity = identity;
     this.pollTask = pollTask;
     this.taskExecutor = taskExecutor;
     this.pollerOptions = pollerOptions;
     this.metricsScope = metricsScope;
     this.pollerAutoScaler = AutoScalerFactory.getInstance().createAutoScaler(pollerOptions);
+    this.threadFactoryWrapper = threadFactoryWrapper;
   }
 
   @Override
@@ -116,8 +122,10 @@ public void start() {
             TimeUnit.SECONDS,
             new ArrayBlockingQueue<>(pollerOptions.getPollThreadCount()));
     pollExecutor.setThreadFactory(
-        new ExecutorThreadFactory(
-            pollerOptions.getPollThreadNamePrefix(), pollerOptions.getUncaughtExceptionHandler()));
+        threadFactoryWrapper.wrap(
+            new ExecutorThreadFactory(
+                pollerOptions.getPollThreadNamePrefix(),
+                pollerOptions.getUncaughtExceptionHandler())));
 
     pollBackoffThrottler =
         new BackoffThrottler(

src/main/java/com/uber/cadence/internal/worker/SingleWorkerOptions.java

@@ -21,6 +21,7 @@
 import com.uber.cadence.converter.DataConverter;
 import com.uber.cadence.converter.JsonDataConverter;
 import com.uber.cadence.internal.metrics.NoopScope;
+import com.uber.cadence.worker.ThreadFactoryWrapper;
 import com.uber.m3.tally.Scope;
 import io.opentracing.Tracer;
 import java.time.Duration;
@@ -46,6 +47,7 @@ public static final class Builder {
     private boolean enableLoggingInReplay;
     private List<ContextPropagator> contextPropagators;
     private Tracer tracer;
+    private ThreadFactoryWrapper threadFactoryWrapper;
 
     private Builder() {}
 
@@ -59,6 +61,7 @@ public Builder(SingleWorkerOptions options) {
       this.enableLoggingInReplay = options.getEnableLoggingInReplay();
       this.contextPropagators = options.getContextPropagators();
       this.tracer = options.getTracer();
+      this.threadFactoryWrapper = options.getThreadFactoryWrapper();
     }
 
     public Builder setIdentity(String identity) {
@@ -107,6 +110,11 @@ public Builder setTracer(Tracer tracer) {
       return this;
     }
 
+    public Builder setThreadFactoryWrapper(ThreadFactoryWrapper threadFactoryWrapper) {
+      this.threadFactoryWrapper = threadFactoryWrapper;
+      return this;
+    }
+
     public SingleWorkerOptions build() {
       if (pollerOptions == null) {
         pollerOptions =
@@ -134,7 +142,8 @@ public SingleWorkerOptions build() {
           metricsScope,
           enableLoggingInReplay,
           contextPropagators,
-          tracer);
+          tracer,
+          threadFactoryWrapper);
     }
   }
 
@@ -147,6 +156,7 @@ public SingleWorkerOptions build() {
   private final boolean enableLoggingInReplay;
   private List<ContextPropagator> contextPropagators;
   private final Tracer tracer;
+  private final ThreadFactoryWrapper threadFactoryWrapper;
 
   private SingleWorkerOptions(
       String identity,
@@ -157,7 +167,8 @@ private SingleWorkerOptions(
       Scope metricsScope,
       boolean enableLoggingInReplay,
       List<ContextPropagator> contextPropagators,
-      Tracer tracer) {
+      Tracer tracer,
+      ThreadFactoryWrapper threadFactoryWrapper) {
     this.identity = identity;
     this.dataConverter = dataConverter;
     this.taskExecutorThreadPoolSize = taskExecutorThreadPoolSize;
@@ -167,6 +178,7 @@ private SingleWorkerOptions(
     this.enableLoggingInReplay = enableLoggingInReplay;
     this.contextPropagators = contextPropagators;
     this.tracer = tracer;
+    this.threadFactoryWrapper = threadFactoryWrapper;
   }
 
   public String getIdentity() {
@@ -204,4 +216,8 @@ public List<ContextPropagator> getContextPropagators() {
   public Tracer getTracer() {
     return tracer;
   }
+
+  public ThreadFactoryWrapper getThreadFactoryWrapper() {
+    return threadFactoryWrapper;
+  }
 }

src/main/java/com/uber/cadence/internal/worker/WorkflowWorker.java

@@ -112,7 +112,8 @@ public void start() {
                   options.getIdentity()),
               pollTaskExecutor,
               options.getPollerOptions(),
-              options.getMetricsScope());
+              options.getMetricsScope(),
+              options.getThreadFactoryWrapper());
       poller.start();
       setPoller(poller);
       options.getMetricsScope().counter(MetricsType.WORKER_START_COUNTER).inc(1);

src/main/java/com/uber/cadence/worker/ShadowingWorker.java

@@ -113,6 +113,8 @@ public ShadowingWorker(
             .setTaskListActivitiesPerSecond(options.getTaskListActivitiesPerSecond())
             .setPollerOptions(options.getActivityPollerOptions())
             .setMetricsScope(metricsScope)
+            .setThreadFactoryWrapper(
+                testOptions.getWorkerFactoryOptions().getThreadFactoryWrapper())
             .build();
     activityWorker =
         new SyncActivityWorker(

src/main/java/com/uber/cadence/worker/ThreadFactoryWrapper.java

@@ -0,0 +1,31 @@
+/*
+ *  Modifications Copyright (c) 2017-2021 Uber Technologies Inc.
+ *  Copyright 2012-2016 Amazon.com, Inc. or its affiliates. All Rights Reserved.
+ *
+ *  Licensed under the Apache License, Version 2.0 (the "License"). You may not
+ *  use this file except in compliance with the License. A copy of the License is
+ *  located at
+ *
+ *  http://aws.amazon.com/apache2.0
+ *
+ *  or in the "license" file accompanying this file. This file is distributed on
+ *  an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
+ *  express or implied. See the License for the specific language governing
+ *  permissions and limitations under the License.
+ */
+package com.uber.cadence.worker;
+
+import java.util.concurrent.ThreadFactory;
+
+public interface ThreadFactoryWrapper {
+  ThreadFactory wrap(ThreadFactory delegate);
+
+  static ThreadFactoryWrapper newDefaultInstance() {
+    return new ThreadFactoryWrapper() {
+      @Override
+      public ThreadFactory wrap(ThreadFactory delegate) {
+        return delegate;
+      }
+    };
+  }
+}

src/main/java/com/uber/cadence/worker/Worker.java

@@ -102,6 +102,7 @@ public final class Worker implements Suspendable {
             .setEnableLoggingInReplay(factoryOptions.isEnableLoggingInReplay())
             .setContextPropagators(contextPropagators)
             .setTracer(options.getTracer())
+            .setThreadFactoryWrapper(factoryOptions.getThreadFactoryWrapper())
             .build();
     activityWorker =
         new SyncActivityWorker(
@@ -117,6 +118,7 @@ public final class Worker implements Suspendable {
             .setEnableLoggingInReplay(factoryOptions.isEnableLoggingInReplay())
             .setContextPropagators(contextPropagators)
             .setTracer(options.getTracer())
+            .setThreadFactoryWrapper(factoryOptions.getThreadFactoryWrapper())
             .build();
     SingleWorkerOptions localActivityOptions =
         SingleWorkerOptions.newBuilder()
@@ -128,6 +130,7 @@ public final class Worker implements Suspendable {
             .setEnableLoggingInReplay(factoryOptions.isEnableLoggingInReplay())
             .setContextPropagators(contextPropagators)
             .setTracer(options.getTracer())
+            .setThreadFactoryWrapper(factoryOptions.getThreadFactoryWrapper())
             .build();
     workflowWorker =
         new SyncWorkflowWorker(

src/main/java/com/uber/cadence/worker/WorkerFactory.java

@@ -101,7 +101,10 @@ public WorkerFactory(WorkflowClient workflowClient, WorkerFactoryOptions factory
             TimeUnit.SECONDS,
             new SynchronousQueue<>());
     workflowThreadPool.setThreadFactory(
-        r -> new Thread(r, "workflow-thread-" + workflowThreadCounter.incrementAndGet()));
+        factoryOptions
+            .getThreadFactoryWrapper()
+            .wrap(
+                r -> new Thread(r, "workflow-thread-" + workflowThreadCounter.incrementAndGet())));
 
     if (this.factoryOptions.isDisableStickyExecution()) {
       return;
@@ -136,7 +139,8 @@ public WorkerFactory(WorkflowClient workflowClient, WorkerFactoryOptions factory
                 .setPollThreadNamePrefix(POLL_THREAD_NAME)
                 .setPollThreadCount(this.factoryOptions.getStickyPollerCount())
                 .build(),
-            stickyScope);
+            stickyScope,
+            factoryOptions.getThreadFactoryWrapper());
   }
 
   /**

src/main/java/com/uber/cadence/worker/WorkerFactoryOptions.java

@@ -48,6 +48,7 @@ public static class Builder {
     private int maxWorkflowThreadCount = DEFAULT_MAX_WORKFLOW_THREAD_COUNT;
     private boolean enableLoggingInReplay;
     private int stickyPollerCount = DEFAULT_STICKY_POLLER_COUNT;
+    private ThreadFactoryWrapper threadFactoryWrapper = ThreadFactoryWrapper.newDefaultInstance();
 
     private Builder() {}
 
@@ -106,14 +107,20 @@ public Builder setEnableLoggingInReplay(boolean enableLoggingInReplay) {
       return this;
     }
 
+    public Builder setThreadFactoryWrapper(ThreadFactoryWrapper executorServiceWrapper) {
+      this.threadFactoryWrapper = executorServiceWrapper;
+      return this;
+    }
+
     public WorkerFactoryOptions build() {
       return new WorkerFactoryOptions(
           disableStickyExecution,
           stickyCacheSize,
           maxWorkflowThreadCount,
           stickyTaskScheduleToStartTimeout,
           stickyPollerCount,
-          enableLoggingInReplay);
+          enableLoggingInReplay,
+          threadFactoryWrapper);
     }
   }
 
@@ -123,14 +130,16 @@ public WorkerFactoryOptions build() {
   private Duration stickyTaskScheduleToStartTimeout;
   private boolean enableLoggingInReplay;
   private int stickyPollerCount;
+  private ThreadFactoryWrapper threadFactoryWrapper;
 
   private WorkerFactoryOptions(
       boolean disableStickyExecution,
       int cacheMaximumSize,
       int maxWorkflowThreadCount,
       Duration stickyTaskScheduleToStartTimeout,
       int stickyPollerCount,
-      boolean enableLoggingInReplay) {
+      boolean enableLoggingInReplay,
+      ThreadFactoryWrapper threadServiceWrapper) {
     Preconditions.checkArgument(cacheMaximumSize > 0, "cacheMaximumSize should be greater than 0");
     Preconditions.checkArgument(
         maxWorkflowThreadCount > 0, "maxWorkflowThreadCount should be greater than 0");
@@ -141,6 +150,7 @@ private WorkerFactoryOptions(
     this.stickyPollerCount = stickyPollerCount;
     this.enableLoggingInReplay = enableLoggingInReplay;
     this.stickyTaskScheduleToStartTimeout = stickyTaskScheduleToStartTimeout;
+    this.threadFactoryWrapper = threadServiceWrapper;
   }
 
   public int getMaxWorkflowThreadCount() {
@@ -166,4 +176,8 @@ public int getStickyPollerCount() {
   public Duration getStickyTaskScheduleToStartTimeout() {
     return stickyTaskScheduleToStartTimeout;
   }
+
+  public ThreadFactoryWrapper getThreadFactoryWrapper() {
+    return threadFactoryWrapper;
+  }
 }