From f384cafed0f7e066ccb8ed7536f5f70abdadc78e Mon Sep 17 00:00:00 2001
From: Aman Mangal <aman@dgraph.io>
Date: Fri, 25 Aug 2023 17:23:13 +0530
Subject: [PATCH] fix(restore): use different map directory for each group
 (#8047) (#8972)

---
 ee/backup/run.go         |  6 +++++-
 worker/online_restore.go | 20 +++++++++++++++-----
 worker/restore_map.go    |  8 ++++----
 3 files changed, 24 insertions(+), 10 deletions(-)

diff --git a/ee/backup/run.go b/ee/backup/run.go
index cdde134e6e7..eb5b666fe71 100644
--- a/ee/backup/run.go
+++ b/ee/backup/run.go
@@ -429,7 +429,11 @@ func runExportBackup() error {
 
 	mapDir, err := os.MkdirTemp(x.WorkerConfig.TmpDir, "restore-export")
 	x.Check(err)
-	defer os.RemoveAll(mapDir)
+	defer func() {
+		if err := os.RemoveAll(mapDir); err != nil {
+			glog.Warningf("Error removing temp restore-export dir: %v", err)
+		}
+	}()
 	glog.Infof("Created temporary map directory: %s\n", mapDir)
 
 	// TODO: Can probably make this procesing concurrent.
diff --git a/worker/online_restore.go b/worker/online_restore.go
index 9b5cdfdb0a2..5672513313b 100644
--- a/worker/online_restore.go
+++ b/worker/online_restore.go
@@ -323,7 +323,11 @@ func handleRestoreProposal(ctx context.Context, req *pb.RestoreRequest, pidx uin
 
 	mapDir, err := os.MkdirTemp(x.WorkerConfig.TmpDir, "restore-map")
 	x.Check(err)
-	defer os.RemoveAll(mapDir)
+	defer func() {
+		if err := os.RemoveAll(mapDir); err != nil {
+			glog.Warningf("Error removing temp restore-map dir: %v", err)
+		}
+	}()
 	glog.Infof("Created temporary map directory: %s\n", mapDir)
 
 	// Map the backup.
@@ -537,10 +541,6 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive,
 		}
 	}
 
-	mapDir, err := os.MkdirTemp(x.WorkerConfig.TmpDir, "restore-map")
-	x.Check(err)
-	defer os.RemoveAll(mapDir)
-
 	for gid := range manifest.Groups {
 		req := &pb.RestoreRequest{
 			Location:          location,
@@ -549,6 +549,16 @@ func RunOfflineRestore(dir, location, backupId, keyFile string, key x.Sensitive,
 			EncryptionKeyFile: keyFile,
 			RestoreTs:         1,
 		}
+		mapDir, err := os.MkdirTemp(x.WorkerConfig.TmpDir, "restore-map")
+		if err != nil {
+			return LoadResult{Err: errors.Wrapf(err, "Failed to create temp map directory")}
+		}
+		defer func() {
+			if err := os.RemoveAll(mapDir); err != nil {
+				glog.Warningf("Error removing temp restore-map dir: %v", err)
+			}
+		}()
+
 		if _, err := RunMapper(req, mapDir); err != nil {
 			return LoadResult{Err: errors.Wrap(err, "RunRestore failed to map")}
 		}
diff --git a/worker/restore_map.go b/worker/restore_map.go
index 30b0fc3fb56..99c32e52739 100644
--- a/worker/restore_map.go
+++ b/worker/restore_map.go
@@ -304,11 +304,11 @@ func (m *mapper) processReqCh(ctx context.Context) error {
 	toBuffer := func(kv *bpb.KV, version uint64) error {
 		key := y.KeyWithTs(kv.Key, version)
 		sz := kv.Size()
-		buf := buf.SliceAllocate(2 + len(key) + sz)
+		b := buf.SliceAllocate(2 + len(key) + sz)
 
-		binary.BigEndian.PutUint16(buf[0:2], uint16(len(key)))
-		x.AssertTrue(copy(buf[2:], key) == len(key))
-		_, err := kv.MarshalToSizedBuffer(buf[2+len(key):])
+		binary.BigEndian.PutUint16(b[0:2], uint16(len(key)))
+		x.AssertTrue(copy(b[2:], key) == len(key))
+		_, err := kv.MarshalToSizedBuffer(b[2+len(key):])
 		return err
 	}