fix(restore): use different map directory for each group (#8047) (#8972)

dgraph-io · Aug 25, 2023 · a668e04 · a668e04
1 parent 69d14a8
commit a668e04
Show file tree

Hide file tree

Showing 3 changed files with 24 additions and 10 deletions.
diff --git 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
@@ -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
@@ -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
 	}