[FLINK-17017][runtime] Allow to disable batch slot request timeout check

shuttie · Jun 11, 2020 · 541f28f · 541f28f
1 parent f6c275e
commit 541f28f
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/slotpool/SlotPool.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/slotpool/SlotPool.java
@@ -186,6 +186,14 @@ CompletableFuture<PhysicalSlot> requestNewAllocatedBatchSlot(
 		@Nonnull SlotRequestId slotRequestId,
 		@Nonnull ResourceProfile resourceProfile);
 
+	/**
+	 * Disables batch slot request timeout check. Invoked when someone else wants to
+	 * take over the timeout check responsibility.
+	 */
+	default void disableBatchSlotRequestTimeoutCheck() {
+		throw new UnsupportedOperationException("Not properly implemented.");
+	}
+
 	/**
 	 * Create report about the allocated slots belonging to the specified task manager.
 	 *

diff --git a/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/slotpool/SlotPoolImpl.java b/flink-runtime/src/main/java/org/apache/flink/runtime/jobmaster/slotpool/SlotPoolImpl.java
@@ -134,6 +134,8 @@ public class SlotPoolImpl implements SlotPool {
 
 	private ComponentMainThreadExecutor componentMainThreadExecutor;
 
+	private boolean batchSlotRequestTimeoutCheckEnabled;
+
 	// ------------------------------------------------------------------------
 
 	public SlotPoolImpl(
@@ -160,6 +162,8 @@ public SlotPoolImpl(
 		this.jobManagerAddress = null;
 
 		this.componentMainThreadExecutor = null;
+
+		this.batchSlotRequestTimeoutCheckEnabled = true;
 	}
 
 	// ------------------------------------------------------------------------
@@ -453,6 +457,11 @@ public CompletableFuture<PhysicalSlot> requestNewAllocatedBatchSlot(
 			.thenApply(Function.identity());
 	}
 
+	@Override
+	public void disableBatchSlotRequestTimeoutCheck() {
+		batchSlotRequestTimeoutCheckEnabled = false;
+	}
+
 	@Override
 	@Nonnull
 	public Collection<SlotInfoWithUtilization> getAvailableSlotsInformation() {
@@ -874,6 +883,10 @@ protected void checkIdleSlot() {
 	}
 
 	protected void checkBatchSlotTimeout() {
+		if (!batchSlotRequestTimeoutCheckEnabled) {
+			return;
+		}
+
 		final Collection<PendingRequest> pendingBatchRequests = getPendingBatchRequests();
 
 		if (!pendingBatchRequests.isEmpty()) {