Clean up unused memory

Author: Tony Cui

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/CancellationSharer.java

@@ -36,7 +36,7 @@
  * It manages the lifecycle of the original attempt and any subsequent hedged attempts.
  */
 class CancellationSharer extends AbstractApiFuture<PublishResponse> {
-  private final Publisher.OutstandingBatch batch;
+  private Publisher.OutstandingBatch batch;
   private final Publisher publisher;
 
   // Guarded by lock
@@ -47,6 +47,11 @@ class CancellationSharer extends AbstractApiFuture<PublishResponse> {
   private final Lock lock = new ReentrantLock();
   private final AtomicBoolean isInQueue = new AtomicBoolean(false);
 
+  private void cleanupLocked() {
+    runningAttempts.clear();
+    this.batch = null;
+  }
+
   CancellationSharer(final Publisher.OutstandingBatch batch, final Publisher publisher) {
     this.batch = batch;
     this.publisher = publisher;
@@ -90,19 +95,18 @@ private void handleAttemptSuccess(final int attemptNumber, final PublishResponse
       batch.successfulAttempt = attemptNumber;
       set(response);
       cancelAllExceptLocked(attemptNumber);
+      cleanupLocked();
     } finally {
       lock.unlock();
     }
     publisher.refillTokenBucket();
   }
 
   private void handleAttemptFailure(final int attemptNumber, final Throwable t) {
-    boolean shouldRemoveFromQueue = false;
     lock.lock();
     try {
       if (done) {
-        return; // <-- Exit early before modifying runningAttempts to avoid
-                // ConcurrentModificationException
+        return;
       }
       runningAttempts.remove(attemptNumber);
       lastError = t;
@@ -117,17 +121,11 @@ private void handleAttemptFailure(final int attemptNumber, final Throwable t) {
         done = true;
         setException(lastError);
         cancelAllLocked();
-        if (isInQueue.get()) {
-          shouldRemoveFromQueue = true;
-        }
+        cleanupLocked();
       }
     } finally {
       lock.unlock();
     }
-
-    if (shouldRemoveFromQueue) {
-      publisher.removeFromHedgingQueue(this);
-    }
   }
 
   void checkCompletionOnQueueExit() {
@@ -139,6 +137,7 @@ void checkCompletionOnQueueExit() {
             lastError != null
                 ? lastError
                 : new RuntimeException("Hedging failed with no active attempts"));
+        cleanupLocked();
       }
     } finally {
       lock.unlock();
@@ -148,24 +147,17 @@ void checkCompletionOnQueueExit() {
   @Override
   public boolean cancel(final boolean mayInterruptIfRunning) {
     boolean cancelled = false;
-    boolean shouldRemoveFromQueue = false;
     lock.lock();
     try {
       if (super.cancel(mayInterruptIfRunning)) {
         cancelled = true;
         done = true;
-        if (isInQueue.get()) {
-          shouldRemoveFromQueue = true;
-        }
         cancelAllLocked();
+        cleanupLocked();
       }
     } finally {
       lock.unlock();
     }
-
-    if (shouldRemoveFromQueue) {
-      publisher.removeFromHedgingQueue(this);
-    }
     return cancelled;
   }
 
@@ -190,7 +182,12 @@ AtomicBoolean isInQueue() {
     return isInQueue;
   }
 
-  Publisher.OutstandingBatch getBatch() {
-    return batch;
+  Publisher.OutstandingBatch getBatchIfActive() {
+    lock.lock();
+    try {
+      return done ? null : batch;
+    } finally {
+      lock.unlock();
+    }
   }
 }

java-pubsub/google-cloud-pubsub/src/main/java/com/google/cloud/pubsub/v1/Publisher.java

@@ -760,21 +760,6 @@ public void run() {
     }
   }
 
-  void removeFromHedgingQueue(CancellationSharer coordinator) {
-    queueLock.lock();
-    try {
-      Iterator<HedgedRequest> iterator = hedgingQueue.iterator();
-      while (iterator.hasNext()) {
-        if (iterator.next().getCoordinator() == coordinator) {
-          iterator.remove();
-          coordinator.isInQueue().set(false);
-        }
-      }
-    } finally {
-      queueLock.unlock();
-    }
-  }
-
   private void processQueue() {
     queueLock.lock();
     try {
@@ -786,7 +771,8 @@ private void processQueue() {
         hedgingQueue.poll();
 
         CancellationSharer coordinator = item.getCoordinator();
-        if (coordinator.isDone()) {
+        OutstandingBatch batch = coordinator.getBatchIfActive();
+        if (batch == null) {
           coordinator.isInQueue().set(false);
           continue;
         }
@@ -799,15 +785,14 @@ private void processQueue() {
           hedgingQueue.add(nextItem);
 
           // Start Hedged Attempt
-          ApiFuture<PublishResponse> hedgedFuture =
-              publishCall(coordinator.getBatch(), item.getAttemptNumber());
+          ApiFuture<PublishResponse> hedgedFuture = publishCall(batch, item.getAttemptNumber());
           coordinator.addAttempt(item.getAttemptNumber(), hedgedFuture);
         } else {
           loggingUtil.logPublisher(
               LoggingUtil.SubSystem.PUBLISH_HEDGED,
               Level.FINER,
               "Hedging rate limited due to lack of tokens.",
-              coordinator.getBatch().getMessageWrappers().get(0));
+              batch.getMessageWrappers().get(0));
           coordinator.isInQueue().set(false);
           coordinator.checkCompletionOnQueueExit();
         }