Skip to content

Commit af7b5b7

Browse files
Fix memory leak of command inputs (#2262)
Fix memory leak of command inputs
1 parent 92692b4 commit af7b5b7

File tree

3 files changed

+32
-11
lines changed

3 files changed

+32
-11
lines changed

temporal-sdk/src/main/java/io/temporal/internal/statemachines/ActivityStateMachine.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,7 @@ public void createScheduleActivityTaskCommand() {
275275
.setCommandType(CommandType.COMMAND_TYPE_SCHEDULE_ACTIVITY_TASK)
276276
.setScheduleActivityTaskCommandAttributes(parameters.getAttributes())
277277
.build());
278+
parameters = null; // avoiding retaining large input for the duration of the activity
278279
}
279280

280281
private void setStartedCommandEventId() {
@@ -455,7 +456,6 @@ private void createRequestCancelActivityTaskCommand() {
455456
RequestCancelActivityTaskCommandAttributes.newBuilder()
456457
.setScheduledEventId(getInitialCommandEventId()))
457458
.build());
458-
parameters = null; // avoid retaining large input for the duration of the activity
459459
}
460460

461461
public static class FailureResult {

temporal-sdk/src/main/java/io/temporal/internal/statemachines/EntityStateMachineInitialCommand.java

Lines changed: 14 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -20,16 +20,17 @@
2020

2121
package io.temporal.internal.statemachines;
2222

23+
import static io.temporal.internal.common.WorkflowExecutionUtils.isCommandEvent;
24+
2325
import io.temporal.api.command.v1.Command;
2426
import io.temporal.api.enums.v1.CommandType;
27+
import io.temporal.api.history.v1.HistoryEvent;
2528
import io.temporal.workflow.Functions;
2629
import javax.annotation.Nullable;
2730

2831
class EntityStateMachineInitialCommand<State, ExplicitEvent, Data>
2932
extends EntityStateMachineBase<State, ExplicitEvent, Data> {
30-
3133
private CancellableCommand command;
32-
3334
private long initialCommandEventId;
3435

3536
public EntityStateMachineInitialCommand(
@@ -56,7 +57,17 @@ protected final void addCommand(Command command) {
5657
}
5758

5859
protected final void cancelCommand() {
59-
command.cancel();
60+
if (command != null) {
61+
command.cancel();
62+
}
63+
}
64+
65+
public WorkflowStateMachines.HandleEventStatus handleEvent(
66+
HistoryEvent event, boolean hasNextEvent) {
67+
if (isCommandEvent(event)) {
68+
command = null;
69+
}
70+
return super.handleEvent(event, hasNextEvent);
6071
}
6172

6273
protected long getInitialCommandEventId() {

temporal-sdk/src/main/java/io/temporal/internal/sync/SyncWorkflowContext.java

Lines changed: 17 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -278,16 +278,20 @@ public <T> ActivityOutput<T> executeActivity(ActivityInput<T> input) {
278278
ActivityOutput<Optional<Payloads>> output =
279279
executeActivityOnce(input.getActivityName(), input.getOptions(), input.getHeader(), args);
280280

281+
// Avoid passing the input to the output handle as it causes the input to be retained for the
282+
// duration of the operation.
283+
Type resultType = input.getResultType();
284+
Class<T> resultClass = input.getResultClass();
281285
return new ActivityOutput<>(
282286
output.getActivityId(),
283287
output
284288
.getResult()
285289
.handle(
286290
(r, f) -> {
287291
if (f == null) {
288-
return input.getResultType() != Void.TYPE
292+
return resultType != Void.TYPE
289293
? dataConverterWithActivityContext.fromPayloads(
290-
0, r, input.getResultClass(), input.getResultType())
294+
0, r, resultClass, resultType)
291295
: null;
292296
} else {
293297
throw dataConverterWithActivityContext.failureToException(
@@ -412,13 +416,16 @@ public <R> LocalActivityOutput<R> executeLocalActivity(LocalActivityInput<R> inp
412416
null,
413417
serializedResult);
414418

419+
// Avoid passing the input to the output handle as it causes the input to be retained for the
420+
// duration of the operation.
421+
Type resultType = input.getResultType();
422+
Class<R> resultClass = input.getResultClass();
415423
Promise<R> result =
416424
serializedResult.handle(
417425
(r, f) -> {
418426
if (f == null) {
419-
return input.getResultClass() != Void.TYPE
420-
? dataConverterWithActivityContext.fromPayloads(
421-
0, r, input.getResultClass(), input.getResultType())
427+
return resultClass != Void.TYPE
428+
? dataConverterWithActivityContext.fromPayloads(0, r, resultClass, resultType)
422429
: null;
423430
} else {
424431
throw dataConverterWithActivityContext.failureToException(
@@ -708,11 +715,14 @@ public <R> ChildWorkflowOutput<R> executeChildWorkflow(ChildWorkflowInput<R> inp
708715
return null;
709716
});
710717

718+
// Avoid passing the input to the output handle as it causes the input to be retained for the
719+
// duration of the operation.
720+
Type resultType = input.getResultType();
721+
Class<R> resultClass = input.getResultClass();
711722
Promise<R> result =
712723
resultPromise.thenApply(
713724
(b) ->
714-
dataConverterWithChildWorkflowContext.fromPayloads(
715-
0, b, input.getResultClass(), input.getResultType()));
725+
dataConverterWithChildWorkflowContext.fromPayloads(0, b, resultClass, resultType));
716726
return new ChildWorkflowOutput<>(result, executionPromise);
717727
}
718728

0 commit comments

Comments
 (0)