From 08c783bcbb442742ae299bd1f2fb15382951eb13 Mon Sep 17 00:00:00 2001
From: akorotkov <akorotkov@aerospike.com>
Date: Thu, 9 Jan 2025 14:52:21 +0200
Subject: [PATCH] fix race condition

---
 pkg/service/estimates.go    |  3 +--
 pkg/service/jobs_holder.go  | 12 ++++++++++++
 pkg/service/restore_data.go |  7 +------
 3 files changed, 14 insertions(+), 8 deletions(-)

diff --git a/pkg/service/estimates.go b/pkg/service/estimates.go
index 0ffe971f..7dd5b3b1 100644
--- a/pkg/service/estimates.go
+++ b/pkg/service/estimates.go
@@ -58,8 +58,7 @@ func RestoreJobStatus(job *jobInfo) *model.RestoreJobStatus {
 	}
 
 	if job.status == model.JobStatusRunning {
-		status.CurrentRestore = NewRunningJob(job.startTime, status.ReadRecords,
-			job.totalRecords)
+		status.CurrentRestore = NewRunningJob(job.startTime, status.ReadRecords, job.totalRecords)
 	}
 
 	if job.err != nil {
diff --git a/pkg/service/jobs_holder.go b/pkg/service/jobs_holder.go
index 88f52d59..b2ccc05e 100644
--- a/pkg/service/jobs_holder.go
+++ b/pkg/service/jobs_holder.go
@@ -74,6 +74,18 @@ func (h *RestoreJobsHolder) finishJob(id model.RestoreJobID, err error) {
 	})
 }
 
+func (h *RestoreJobsHolder) getStatus(id model.RestoreJobID) (*model.RestoreJobStatus, error) {
+	var result *model.RestoreJobStatus
+	h.Apply(id, func(value *jobInfo) {
+		result = RestoreJobStatus(value)
+	})
+
+	if result != nil {
+		return result, nil
+	}
+	return nil, NewErrJobNotFound(id)
+}
+
 func (h *RestoreJobsHolder) getJob(id model.RestoreJobID) (*jobInfo, error) {
 	if job, exists := h.Load(id); exists {
 		return job, nil
diff --git a/pkg/service/restore_data.go b/pkg/service/restore_data.go
index d548243c..164a01ef 100644
--- a/pkg/service/restore_data.go
+++ b/pkg/service/restore_data.go
@@ -247,12 +247,7 @@ func (r *dataRestorer) restoreFromPath(
 
 // JobStatus returns the status of the job with the given id.
 func (r *dataRestorer) JobStatus(jobID model.RestoreJobID) (*model.RestoreJobStatus, error) {
-	job, err := r.restoreJobs.getJob(jobID)
-	if err != nil {
-		return nil, err
-	}
-
-	return RestoreJobStatus(job), nil
+	return r.restoreJobs.getStatus(jobID)
 }
 
 func recordsInBackup(ctx context.Context, request *model.RestoreRequest) (uint64, error) {