From 68bcecd7d41c3118cc6034681fdf2a2c649d6726 Mon Sep 17 00:00:00 2001
From: Meng Xu <meng_xu@apple.com>
Date: Mon, 4 May 2020 11:31:39 -0700
Subject: [PATCH] FastRestoreLoader:Clean batchData when a version batch
 finishes

Add sanity check for each loader phase to ensure an older version batch
will not be executed once a version batch has been marked as finished.
---
 fdbserver/RestoreLoader.actor.cpp     | 7 +++++++
 fdbserver/RestoreLoader.actor.h       | 1 +
 fdbserver/RestoreRoleCommon.actor.cpp | 3 +++
 3 files changed, 11 insertions(+)

diff --git a/fdbserver/RestoreLoader.actor.cpp b/fdbserver/RestoreLoader.actor.cpp
index 3abfa1f060..d6f6379859 100644
--- a/fdbserver/RestoreLoader.actor.cpp
+++ b/fdbserver/RestoreLoader.actor.cpp
@@ -336,6 +336,8 @@ ACTOR Future<Void> handleLoadFileRequest(RestoreLoadFileRequest req, Reference<R
 	    .detail("NotProcessed", !paramExist)
 	    .detail("Processed", isReady)
 	    .detail("CurrentMemory", getSystemStatistics().processMemory);
+	// Loader destroy batchData once the batch finishes and self->finishedBatch.set(req.batchIndex);
+	ASSERT(self->finishedBatch.get() < req.batchIndex);
 
 	wait(isSchedulable(self, req.batchIndex, __FUNCTION__));
 
@@ -376,6 +378,8 @@ ACTOR Future<Void> handleSendMutationsRequest(RestoreSendMutationsToAppliersRequ
 	    .detail("BatchIndex", req.batchIndex)
 	    .detail("UseRangeFile", req.useRangeFile)
 	    .detail("LoaderSendStatus", batchStatus->toString());
+	// Loader destroy batchData once the batch finishes and self->finishedBatch.set(req.batchIndex);
+	ASSERT(self->finishedBatch.get() < req.batchIndex);
 
 	// Ensure each file is sent exactly once by using batchStatus->sendAllLogs and batchStatus->sendAllRanges
 	if (!req.useRangeFile) {
@@ -945,6 +949,9 @@ ACTOR Future<Void> handleFinishVersionBatchRequest(RestoreVersionBatchRequest re
 	wait(self->finishedBatch.whenAtLeast(req.batchIndex - 1));
 	if (self->finishedBatch.get() == req.batchIndex - 1) {
 		self->finishedBatch.set(req.batchIndex);
+		// Clean up batchData
+		self->batch.erase(req.batchIndex);
+		self->status.erase(req.batchIndex);
 	}
 	if (self->delayedActors > 0) {
 		self->checkMemory.trigger();
diff --git a/fdbserver/RestoreLoader.actor.h b/fdbserver/RestoreLoader.actor.h
index d2469fbcaa..f6fc772819 100644
--- a/fdbserver/RestoreLoader.actor.h
+++ b/fdbserver/RestoreLoader.actor.h
@@ -177,6 +177,7 @@ struct RestoreLoaderData : RestoreRoleData, public ReferenceCounted<RestoreLoade
 	void resetPerRestoreRequest() {
 		batch.clear();
 		status.clear();
+		finishedBatch = NotifiedVersion(0);
 	}
 
 	void initBackupContainer(Key url) {
diff --git a/fdbserver/RestoreRoleCommon.actor.cpp b/fdbserver/RestoreRoleCommon.actor.cpp
index f8c75c8524..281a8ec37e 100644
--- a/fdbserver/RestoreRoleCommon.actor.cpp
+++ b/fdbserver/RestoreRoleCommon.actor.cpp
@@ -57,6 +57,9 @@ ACTOR Future<Void> handleInitVersionBatchRequest(RestoreVersionBatchRequest req,
 	    .detail("BatchIndex", req.batchIndex)
 	    .detail("Role", getRoleStr(self->role))
 	    .detail("VersionBatchNotifiedVersion", self->versionBatchId.get());
+	// Loader destroy batchData once the batch finishes and self->finishedBatch.set(req.batchIndex);
+	ASSERT(self->finishedBatch.get() < req.batchIndex);
+
 	// batchId is continuous. (req.batchIndex-1) is the id of the just finished batch.
 	wait(self->versionBatchId.whenAtLeast(req.batchIndex - 1));