[BEAM-11417] StreamingDataflowWorker can leak UnboundedSource finalization callbacks

• Type: Bug
• Status: Resolved
• Priority: P1
• Resolution: Fixed
• Fix Version/s: 2.27.0
• Component/s: runner-dataflow
• Assignee: Boyuan Zhang
• Reporter: Daniel Mills
• Created: 08/Dec/20 01:39
• Updated: 08/Dec/20 20:25
• Resolved: 08/Dec/20 20:25

Description

StreamingDataflowWorker keeps a map of finalization callbacks (https://github.com/apache/beam/blob/master/runners/google-cloud-dataflow-java/worker/src/main/java/org/apache/beam/runners/dataflow/worker/StreamingDataflowWorker.java#L401).  If the Dataflow service loses a callback ID (due to autoscaling etc; they are best-effort), the callback will stay around forever.

This can cause a relatively rapid memory leak for sources like KafkaIO where the callback (the KafkaCheckpointMark) has a reference to the KafkaUnboundedReader object, which keeps a KafkaConsumer object alive.

A simple fix would be to change the ConcurrentHashMap to a guava Cache with a timeout on elements.

Issue Links

• links to: GitHub Pull Request #13507