Skip to content

Commit 1fb844a

Browse files
mar-kolyamar-kolya
committed
Make CommonTaskExecutor periodic tasks safe
* Verify that we can schedule task and catch exceptions. This should help to avoid additional exceptions on app crash during startup. * Avoid holding strong references from within executor to make sure that things can get GCed.
1 parent af188c2 commit 1fb844a

4 files changed

Lines changed: 169 additions & 50 deletions

File tree

‎dd-java-agent/agent-tooling/src/main/java/datadog/trace/agent/tooling/Cleaner.java‎

Lines changed: 11 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
package datadog.trace.agent.tooling;
22

33
import datadog.common.exec.CommonTaskExecutor;
4-
import java.lang.ref.WeakReference;
5-
import java.util.concurrent.RejectedExecutionException;
6-
import java.util.concurrent.ScheduledFuture;
4+
import datadog.common.exec.CommonTaskExecutor.Task;
75
import java.util.concurrent.TimeUnit;
86
import lombok.extern.slf4j.Slf4j;
97

@@ -12,47 +10,25 @@ class Cleaner {
1210

1311
<T> void scheduleCleaning(
1412
final T target, final Adapter<T> adapter, final long frequency, final TimeUnit unit) {
15-
final CleanupRunnable<T> command = new CleanupRunnable<>(target, adapter);
16-
if (CommonTaskExecutor.INSTANCE.isShutdown()) {
17-
log.warn(
18-
"Cleaning scheduled but task scheduler is shutdown. Target won't be cleaned {}", target);
19-
} else {
20-
try {
21-
// Schedule job and save future to allow job to be canceled if target is GC'd.
22-
command.setFuture(
23-
CommonTaskExecutor.INSTANCE.scheduleAtFixedRate(command, frequency, frequency, unit));
24-
} catch (final RejectedExecutionException e) {
25-
log.warn("Cleaning task rejected. Target won't be cleaned {}", target);
26-
}
27-
}
28-
}
29-
30-
public interface Adapter<T> {
31-
void clean(T target);
13+
CommonTaskExecutor.INSTANCE.scheduleAtFixedRate(
14+
new CleaningTask(adapter), target, frequency, frequency, unit, "cleaner for " + target);
3215
}
3316

34-
private static class CleanupRunnable<T> implements Runnable {
35-
private final WeakReference<T> target;
17+
// Important to use explicit class to avoid implicit hard references to target
18+
private static class CleaningTask<T> implements Task<T> {
3619
private final Adapter<T> adapter;
37-
private volatile ScheduledFuture<?> future = null;
3820

39-
private CleanupRunnable(final T target, final Adapter<T> adapter) {
40-
this.target = new WeakReference<>(target);
21+
public CleaningTask(final Adapter<T> adapter) {
4122
this.adapter = adapter;
4223
}
4324

4425
@Override
45-
public void run() {
46-
final T t = target.get();
47-
if (t != null) {
48-
adapter.clean(t);
49-
} else if (future != null) {
50-
future.cancel(false);
51-
}
26+
public void run(final T target) {
27+
adapter.clean(target);
5228
}
29+
}
5330

54-
public void setFuture(final ScheduledFuture<?> future) {
55-
this.future = future;
56-
}
31+
public interface Adapter<T> {
32+
void clean(T target);
5733
}
5834
}

‎dd-trace-ot/src/main/java/datadog/opentracing/PendingTrace.java‎

Lines changed: 21 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
package datadog.opentracing;
22

33
import datadog.common.exec.CommonTaskExecutor;
4+
import datadog.common.exec.CommonTaskExecutor.Task;
45
import datadog.opentracing.scopemanager.ContinuableScope;
56
import datadog.trace.common.util.Clock;
67
import java.io.Closeable;
@@ -296,7 +297,13 @@ private static class SpanCleaner implements Runnable, Closeable {
296297
Collections.newSetFromMap(new ConcurrentHashMap<PendingTrace, Boolean>());
297298

298299
public SpanCleaner() {
299-
CommonTaskExecutor.INSTANCE.scheduleAtFixedRate(this, 0, CLEAN_FREQUENCY, TimeUnit.SECONDS);
300+
CommonTaskExecutor.INSTANCE.scheduleAtFixedRate(
301+
SpanCleanerTask.INSTANCE,
302+
this,
303+
0,
304+
CLEAN_FREQUENCY,
305+
TimeUnit.SECONDS,
306+
"Pending trace cleaner");
300307
}
301308

302309
@Override
@@ -312,4 +319,17 @@ public void close() {
312319
run();
313320
}
314321
}
322+
323+
/*
324+
* Important to use explicit class to avoid implicit hard references to cleaners from within executor.
325+
*/
326+
private static class SpanCleanerTask implements Task<SpanCleaner> {
327+
328+
static final SpanCleanerTask INSTANCE = new SpanCleanerTask();
329+
330+
@Override
331+
public void run(final SpanCleaner target) {
332+
target.run();
333+
}
334+
}
315335
}

‎dd-trace-ot/src/main/java/datadog/trace/common/writer/ddagent/BatchWritingDisruptor.java‎

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
import com.lmax.disruptor.EventHandler;
44
import datadog.common.exec.CommonTaskExecutor;
5+
import datadog.common.exec.CommonTaskExecutor.Task;
56
import datadog.common.exec.DaemonThreadFactory;
67
import datadog.trace.common.writer.DDAgentWriter;
78
import java.util.ArrayList;
@@ -32,17 +33,8 @@ public BatchWritingDisruptor(
3233

3334
if (0 < flushFrequencySeconds) {
3435
// This provides a steady stream of events to enable flushing with a low throughput.
35-
final Runnable heartbeat =
36-
new Runnable() {
37-
@Override
38-
public void run() {
39-
// Only add if the buffer is empty.
40-
if (running && getCurrentCount() == 0) {
41-
disruptor.getRingBuffer().tryPublishEvent(heartbeatTranslator);
42-
}
43-
}
44-
};
45-
CommonTaskExecutor.INSTANCE.scheduleAtFixedRate(heartbeat, 100, 100, TimeUnit.MILLISECONDS);
36+
CommonTaskExecutor.INSTANCE.scheduleAtFixedRate(
37+
new HeartbeatTask(), this, 100, 100, TimeUnit.MILLISECONDS, "disruptor heartbeat");
4638
}
4739
}
4840

@@ -58,6 +50,12 @@ public boolean publish(final byte[] data, final int representativeCount) {
5850
return true;
5951
}
6052

53+
private void heartbeat() {
54+
if (running && getCurrentCount() == 0) {
55+
disruptor.getRingBuffer().tryPublishEvent(heartbeatTranslator);
56+
}
57+
}
58+
6159
// Intentionally not thread safe.
6260
private static class BatchWritingHandler implements EventHandler<DisruptorEvent<byte[]>> {
6361

@@ -162,4 +160,12 @@ private void scheduleNextFlush() {
162160
}
163161
}
164162
}
163+
164+
// Important to use explicit class to avoid implicit hard references to BatchWritingDisruptor
165+
private static final class HeartbeatTask implements Task<BatchWritingDisruptor> {
166+
@Override
167+
public void run(final BatchWritingDisruptor target) {
168+
target.heartbeat();
169+
}
170+
}
165171
}

‎utils/thread-utils/src/main/java/datadog/common/exec/CommonTaskExecutor.java‎

Lines changed: 120 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -1,8 +1,11 @@
11
package datadog.common.exec;
22

3+
import java.lang.ref.WeakReference;
34
import java.util.List;
45
import java.util.concurrent.AbstractExecutorService;
6+
import java.util.concurrent.Delayed;
57
import java.util.concurrent.Executors;
8+
import java.util.concurrent.RejectedExecutionException;
69
import java.util.concurrent.ScheduledExecutorService;
710
import java.util.concurrent.ScheduledFuture;
811
import java.util.concurrent.TimeUnit;
@@ -25,9 +28,48 @@ private CommonTaskExecutor() {
2528
}
2629
}
2730

28-
public ScheduledFuture<?> scheduleAtFixedRate(
29-
final Runnable command, final long initialDelay, final long period, final TimeUnit unit) {
30-
return executorService.scheduleAtFixedRate(command, initialDelay, period, unit);
31+
/**
32+
* Run {@code task} periodically providing it with {@code target}
33+
*
34+
* <p>Important implementation detail here is that internally we do not hold any strong references
35+
* to {@code target} which means it can be GCed even while periodic task is still scheduled.
36+
*
37+
* <p>If {@code target} is GCed periodic task is canceled.
38+
*
39+
* @param task task to run. Important: must not hold any strong references to target (or anything
40+
* else non static)
41+
* @param target target object to pass to task
42+
* @param initialDelay initialDelay, see {@link
43+
* ScheduledExecutorService#scheduleAtFixedRate(Runnable, long, long, TimeUnit)}
44+
* @param period period, see {@link ScheduledExecutorService#scheduleAtFixedRate(Runnable, long,
45+
* long, TimeUnit)}
46+
* @param unit unit, see {@link ScheduledExecutorService#scheduleAtFixedRate(Runnable, long, long,
47+
* TimeUnit)}
48+
* @param name name to use in logs when task cannot be scheduled
49+
* @return future that can be canceled
50+
*/
51+
public <T> ScheduledFuture<?> scheduleAtFixedRate(
52+
final Task<T> task,
53+
final T target,
54+
final long initialDelay,
55+
final long period,
56+
final TimeUnit unit,
57+
final String name) {
58+
if (CommonTaskExecutor.INSTANCE.isShutdown()) {
59+
log.warn("Periodic task scheduler is shutdown. Will not run: {}", name);
60+
} else {
61+
try {
62+
final PeriodicTask<T> periodicTask = new PeriodicTask<>(task, target);
63+
final ScheduledFuture<?> future =
64+
executorService.scheduleAtFixedRate(
65+
new PeriodicTask<>(task, target), initialDelay, period, unit);
66+
periodicTask.setFuture(future);
67+
return future;
68+
} catch (final RejectedExecutionException e) {
69+
log.warn("Cleaning task rejected. Will not run: {}", name);
70+
}
71+
}
72+
return new UnscheduledFuture(name);
3173
}
3274

3375
@Override
@@ -82,4 +124,79 @@ public void run() {
82124
}
83125
}
84126
}
127+
128+
public interface Task<T> {
129+
void run(T target);
130+
}
131+
132+
public static class PeriodicTask<T> implements Runnable {
133+
private final WeakReference<T> target;
134+
private final Task<T> task;
135+
private volatile ScheduledFuture<?> future = null;
136+
137+
private PeriodicTask(final Task<T> task, final T target) {
138+
this.target = new WeakReference<>(target);
139+
this.task = task;
140+
}
141+
142+
@Override
143+
public void run() {
144+
final T t = target.get();
145+
if (t != null) {
146+
task.run(t);
147+
} else if (future != null) {
148+
future.cancel(false);
149+
}
150+
}
151+
152+
public void setFuture(final ScheduledFuture<?> future) {
153+
this.future = future;
154+
}
155+
}
156+
157+
// Unscheduled future
158+
@Slf4j
159+
public static class UnscheduledFuture implements ScheduledFuture<Object> {
160+
private final String name;
161+
162+
public UnscheduledFuture(final String name) {
163+
this.name = name;
164+
}
165+
166+
@Override
167+
public long getDelay(final TimeUnit unit) {
168+
return 0;
169+
}
170+
171+
@Override
172+
public int compareTo(final Delayed o) {
173+
return 0;
174+
}
175+
176+
@Override
177+
public boolean cancel(final boolean mayInterruptIfRunning) {
178+
log.debug("Cancelling future for: {}", name);
179+
return false;
180+
}
181+
182+
@Override
183+
public boolean isCancelled() {
184+
return false;
185+
}
186+
187+
@Override
188+
public boolean isDone() {
189+
return false;
190+
}
191+
192+
@Override
193+
public Object get() {
194+
return null;
195+
}
196+
197+
@Override
198+
public Object get(final long timeout, final TimeUnit unit) {
199+
return null;
200+
}
201+
}
85202
}

0 commit comments

Comments
 (0)