Async activity inputs potential memory leak

[hansen-jake]

When starting multiple async activities within a workflow, references to the activity inputs are never null’d out so the inputs are never garbage collected, leading to a memory leak. This has caused our workers to run out of memory.

Expected Behavior

The expectation/assumption is that as activities within a workflow complete, references to their inputs are null’d out, making them eligible for garbage collection.

Actual Behavior

After running the below code for a while, a heap dump shows that a strong reference to each activity input is held by the class WorkflowOutboundCallsInterceptor$ActivityInput, even after each activity completes, causing a memory leak.

Steps to Reproduce the Problem

See below code for reproduction.

Specifications

• Version: Java SDK 1.25.1

// NOT A CONTRIBUTION

public class Main {
    public static void main(String[] args) {
        WorkflowServiceStubs service = WorkflowServiceStubs.newLocalServiceStubs();
        WorkflowClient client = WorkflowClient.newInstance(service);
        WorkerFactory factory = WorkerFactory.newInstance(client);
        Worker worker = factory.newWorker("task_queue");
        worker.registerWorkflowImplementationTypes(EventLoopWorkflowImpl.class);
        worker.registerActivitiesImplementations(new ProcessingActivityImpl(client));
        factory.start();

        EventLoopWorkflow workflow = client.newWorkflowStub(
                EventLoopWorkflow.class,
                WorkflowOptions.newBuilder()
                        .setWorkflowId(UUID.randomUUID().toString())
                        .setTaskQueue("task_queue")
                        .build());

        workflow.startEventLoop();
    }

    @WorkflowInterface
    public interface EventLoopWorkflow {

        @WorkflowMethod
        void startEventLoop();

        @SignalMethod
        void addBatch(List<String> batch);
    }

    public static class EventLoopWorkflowImpl implements EventLoopWorkflow {

        private final ConcurrentLinkedQueue<Supplier<?>> eventQueue = new ConcurrentLinkedQueue<>();

        private final ProcessingActivity activity = Workflow.newActivityStub(
                ProcessingActivity.class,
                ActivityOptions.newBuilder()
                        .setStartToCloseTimeout(Duration.ofHours(24))
                        .build());

        @Override
        public void startEventLoop() {
            var getBatches = Async.procedure(activity::getBatches);

            do {
                Workflow.await(() -> getBatches.isCompleted() || !eventQueue.isEmpty());

                Supplier<?> event;
                while ((event = eventQueue.poll()) != null) {
                    event.get();
                }

            } while (!getBatches.isCompleted());
        }

        @Override
        public void addBatch(List<String> batch) {
            eventQueue.add(() -> Async.procedure(activity::processBatch, batch));
        }
    }

    @ActivityInterface
    public interface ProcessingActivity {

        void getBatches();

        void processBatch(List<String> items);
    }

    public static class ProcessingActivityImpl implements ProcessingActivity {

        private final WorkflowClient workflowClient;

        public ProcessingActivityImpl(WorkflowClient workflowClient) {
            this.workflowClient = workflowClient;
        }


        @Override
        public void getBatches() {
            var ctx = Activity.getExecutionContext();
            var workflow = workflowClient.newWorkflowStub(EventLoopWorkflow.class, ctx.getInfo().getWorkflowId());
            int batchId = 0;

            while (true) {
                try {
                    workflow.addBatch(List.of(String.valueOf(batchId++)));
                    Thread.sleep(1000);
                } catch (InterruptedException e) {
                    throw new RuntimeException(e);
                }
            }
        }

        @Override
        public void processBatch(List<String> items) {
            System.out.println("Processing batch " + items.get(0));
        }
    }
}

[Quinn-With-Two-Ns]

Thank you for the detailed reproduction. Running the provided code and looking at a memory snapshot it looks like the activity input is being held because the internal ActivityStateMachine holds it and the state machine is being held by the cancellation callback in the CancellationScope. We could make the cancellation callback hold a weak reference to the state machine or remove the cancellation callback once the ActivityStateMachine is in an terminal state.

[hansen-jake]

What are the tradeoffs of doing one over the other?

[Quinn-With-Two-Ns]

From the user perspective very little. Either would remove the reference to the StateMachine and let the input be GC'd

[hansen-jake]

Is there a target version where this fix can be included?

[Quinn-With-Two-Ns]

Yes I hope to have some improvements here for the next minor SDK release .