From d410d20ffdc9fdb1cee5d97300b1cb7fab004c49 Mon Sep 17 00:00:00 2001
From: Sam Whittle <samuelw@google.com>
Date: Wed, 30 Oct 2024 10:57:34 +0100
Subject: [PATCH] Changes to reduce memory pinned while iterating through state
 backed iterable: - remove reference to completed encoded input page from
 decoder once we have read it. - re-read from cache after loading the next
 page to give eviction a chance to remove blocks

---
 .../fn/data/BeamFnDataGrpcMultiplexer.java    | 58 +++++++++----------
 1 file changed, 29 insertions(+), 29 deletions(-)

diff --git a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java
index 7068664fd470..aa0dea80b0a1 100644
--- a/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java
+++ b/sdks/java/core/src/main/java/org/apache/beam/sdk/fn/data/BeamFnDataGrpcMultiplexer.java
@@ -62,7 +62,6 @@ public class BeamFnDataGrpcMultiplexer implements AutoCloseable {
       receivers;
   private final Cache</*instructionId=*/ String, /*unused=*/ Boolean> poisonedInstructionIds;
 
-
   private static class PoisonedException extends RuntimeException {
     public PoisonedException() {
       super("Instruction poisoned");
@@ -107,36 +106,36 @@ public StreamObserver<BeamFnApi.Elements> getOutboundObserver() {
    * instruction ids ensuring that the receiver will only see {@link BeamFnApi.Elements} with a
    * single instruction id.
    *
-   * <p>The caller must either {@link #unregisterConsumer unregister the consumer} when all messages have
-   * been processed or {@link #poisonInstructionId(String) poison the instruction} if messages for the instruction
-   * should be dropped.
+   * <p>The caller must either {@link #unregisterConsumer unregister the consumer} when all messages
+   * have been processed or {@link #poisonInstructionId(String) poison the instruction} if messages
+   * for the instruction should be dropped.
    */
   public void registerConsumer(
       String instructionId, CloseableFnDataReceiver<BeamFnApi.Elements> receiver) {
     receivers.compute(
-            instructionId,
-            (unused, existing) -> {
-              if (existing != null) {
-                if (!existing.complete(receiver)) {
-                  throw new IllegalArgumentException("Instruction id was registered twice");
-                }
-                return existing;
-              }
-              if (poisonedInstructionIds.getIfPresent(instructionId) != null) {
-                throw new IllegalArgumentException("Instruction id was poisoned");
-              }
-              return CompletableFuture.completedFuture(receiver);
-            });
+        instructionId,
+        (unused, existing) -> {
+          if (existing != null) {
+            if (!existing.complete(receiver)) {
+              throw new IllegalArgumentException("Instruction id was registered twice");
+            }
+            return existing;
+          }
+          if (poisonedInstructionIds.getIfPresent(instructionId) != null) {
+            throw new IllegalArgumentException("Instruction id was poisoned");
+          }
+          return CompletableFuture.completedFuture(receiver);
+        });
   }
 
-  /**
-   * Unregisters a previously registered consumer.
-   * */
+  /** Unregisters a previously registered consumer. */
   public void unregisterConsumer(String instructionId) {
-    @Nullable CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>> receiverFuture =
+    @Nullable
+    CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>> receiverFuture =
         receivers.remove(instructionId);
     if (receiverFuture != null && !receiverFuture.isDone()) {
-      // The future must have been inserted by the inbound observer since registerConsumer completes the future.
+      // The future must have been inserted by the inbound observer since registerConsumer completes
+      // the future.
       throw new IllegalArgumentException("Unregistering consumer which was not registered.");
     }
   }
@@ -144,9 +143,9 @@ public void unregisterConsumer(String instructionId) {
   /**
    * Poisons an instruction id.
    *
-   * <p> Any records for the instruction on the inbound observer will be dropped for the next
-   * {@link #POISONED_INSTRUCTION_ID_CACHE_TIMEOUT}.
-   * */
+   * <p>Any records for the instruction on the inbound observer will be dropped for the next {@link
+   * #POISONED_INSTRUCTION_ID_CACHE_TIMEOUT}.
+   */
   public void poisonInstructionId(String instructionId) {
     poisonedInstructionIds.put(instructionId, Boolean.TRUE);
     @Nullable
@@ -266,16 +265,17 @@ public void onNext(BeamFnApi.Elements value) {
     private void forwardToConsumerForInstructionId(String instructionId, BeamFnApi.Elements value) {
       CloseableFnDataReceiver<BeamFnApi.Elements> consumer;
       try {
-        CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>> consumerFuture = receivers.computeIfAbsent(
+        CompletableFuture<CloseableFnDataReceiver<BeamFnApi.Elements>> consumerFuture =
+            receivers.computeIfAbsent(
                 instructionId,
                 (unused) -> {
                   if (poisonedInstructionIds.getIfPresent(instructionId) != null) {
                     throw new PoisonedException();
                   }
                   LOG.debug(
-                          "Received data for instruction {} without consumer ready. "
-                                  + "Waiting for consumer to be registered.",
-                          instructionId);
+                      "Received data for instruction {} without consumer ready. "
+                          + "Waiting for consumer to be registered.",
+                      instructionId);
                   return new CompletableFuture<>();
                 });
         // The consumer may not be registered until the bundle processor is fully constructed so we