From 06495b90ae873e01fcac8d8eb2a0ce320d908aa9 Mon Sep 17 00:00:00 2001
From: Meng Xu <meng_xu@apple.com>
Date: Wed, 26 Feb 2020 14:35:03 -0800
Subject: [PATCH] FastRestore:Loader:Use isSchedulable to guard OOM

And trigger delayed actors that are blocked on memory to recheck memory.
---
 fdbserver/RestoreApplier.actor.cpp |  4 ++++
 fdbserver/RestoreLoader.actor.cpp  | 11 ++++++++++-
 2 files changed, 14 insertions(+), 1 deletion(-)

diff --git a/fdbserver/RestoreApplier.actor.cpp b/fdbserver/RestoreApplier.actor.cpp
index b4d936411c..891409c98c 100644
--- a/fdbserver/RestoreApplier.actor.cpp
+++ b/fdbserver/RestoreApplier.actor.cpp
@@ -457,6 +457,10 @@ ACTOR static Future<Void> handleApplyToDBRequest(RestoreVersionBatchRequest req,
 			self->finishedBatch.set(req.batchIndex);
 		}
 	}
+
+	if (self->delayedActors > 0) {
+		self->checkMemory.trigger();
+	}
 	req.reply.send(RestoreCommonReply(self->id(), isDuplicated));
 
 	return Void();
diff --git a/fdbserver/RestoreLoader.actor.cpp b/fdbserver/RestoreLoader.actor.cpp
index 8f956bd426..7191f5496d 100644
--- a/fdbserver/RestoreLoader.actor.cpp
+++ b/fdbserver/RestoreLoader.actor.cpp
@@ -201,7 +201,11 @@ ACTOR Future<Void> handleLoadFileRequest(RestoreLoadFileRequest req, Reference<R
 	    .detail("BatchIndex", req.batchIndex)
 	    .detail("ProcessLoadParam", req.param.toString())
 	    .detail("NotProcessed", !paramExist)
-	    .detail("Processed", isReady);
+	    .detail("Processed", isReady)
+		.detail("CurrentMemory", getSystemStatistics().processMemory);
+
+	wait(isSchedulable(self, req.batchIndex, __FUNCTION__);
+
 	if (batchData->processedFileParams.find(req.param) == batchData->processedFileParams.end()) {
 		TraceEvent("FastRestoreLoadFile", self->id())
 		    .detail("BatchIndex", req.batchIndex)
@@ -226,6 +230,8 @@ ACTOR Future<Void> handleLoadFileRequest(RestoreLoadFileRequest req, Reference<R
 	return Void();
 }
 
+// Send buffered mutations to appliers.
+// Do not need to block on low memory usage because this actor should not increase memory usage.
 ACTOR Future<Void> handleSendMutationsRequest(RestoreSendMutationsToAppliersRequest req,
                                               Reference<RestoreLoaderData> self) {
 	state Reference<LoaderBatchData> batchData = self->batch[req.batchIndex];
@@ -745,6 +751,9 @@ ACTOR Future<Void> handleFinishVersionBatchRequest(RestoreVersionBatchRequest re
 	if (self->finishedBatch.get() == req.batchIndex - 1) {
 		self->finishedBatch.set(req.batchIndex);
 	}
+	if (self->delayedActors > 0) {
+		self->checkMemory.trigger();
+	}
 	req.reply.send(RestoreCommonReply(self->id(), false));
 	return Void();
 }
\ No newline at end of file