Cancel receiving mailbox in release cal for in-mem service

ankitsultana · ankitsultana · commit 35b3bad82d17 · 2023-03-09T17:20:31.000-06:00
diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryMailboxService.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/mailbox/InMemoryMailboxService.java
@@ -106,6 +106,10 @@ public ReceivingMailbox<TransferableBlock> getReceivingMailbox(MailboxIdentifier
 
   @Override
   public void releaseReceivingMailbox(MailboxIdentifier mailboxId) {
-    _receivingMailboxCache.invalidate(mailboxId.toString());
+    InMemoryReceivingMailbox receivingMailbox = _receivingMailboxCache.getIfPresent(mailboxId.toString());
+    if (receivingMailbox != null) {
+      receivingMailbox.cancel();
+      _receivingMailboxCache.invalidate(mailboxId.toString());
+    }
   }
 }

Original file line number	Diff line number	Diff line change
`@@ -106,6 +106,10 @@ public ReceivingMailbox<TransferableBlock> getReceivingMailbox(MailboxIdentifier`
`106`	`106`
`107`	`107`	`@Override`
`108`	`108`	`public void releaseReceivingMailbox(MailboxIdentifier mailboxId) {`
`109`		`- _receivingMailboxCache.invalidate(mailboxId.toString());`
	`109`	`+ InMemoryReceivingMailbox receivingMailbox = _receivingMailboxCache.getIfPresent(mailboxId.toString());`
	`110`	`+ if (receivingMailbox != null) {`
	`111`	`+ receivingMailbox.cancel();`
	`112`	`+ _receivingMailboxCache.invalidate(mailboxId.toString());`
	`113`	`+ }`
`110`	`114`	`}`
`111`	`115`	`}`