Fix Services & Update the middleware && others

internal/task/task.go

@@ -1,4 +1,3 @@
-// Filename: internal/task/task.go
 package task
 
 import (
@@ -13,7 +12,9 @@ import (
 )
 
 const (
-	ResultTTL = 60 * time.Minute
+	ResultTTL      = 60 * time.Minute
+	DefaultTimeout = 24 * time.Hour
+	LockTTL        = 30 * time.Minute
 )
 
 type Reporter interface {
@@ -65,14 +66,21 @@ func (s *Task) getIsRunningFlagKey(taskID string) string {
 
 func (s *Task) StartTask(ctx context.Context, keyGroupID uint, taskType, resourceID string, total int, timeout time.Duration) (*Status, error) {
 	lockKey := s.getResourceLockKey(resourceID)
+	taskID := fmt.Sprintf("%d-%d", time.Now().UnixNano(), keyGroupID)
 
-	if existingTaskID, err := s.store.Get(ctx, lockKey); err == nil && len(existingTaskID) > 0 {
+	locked, err := s.store.SetNX(ctx, lockKey, []byte(taskID), LockTTL)
+	if err != nil {
+		return nil, fmt.Errorf("failed to acquire task lock: %w", err)
+	}
+	if !locked {
+		existingTaskID, _ := s.store.Get(ctx, lockKey)
 		return nil, fmt.Errorf("a task is already running for this resource (ID: %s)", string(existingTaskID))
 	}
 
-	taskID := fmt.Sprintf("%d-%d", time.Now().UnixNano(), keyGroupID)
-	taskKey := s.getTaskDataKey(taskID)
-	runningFlagKey := s.getIsRunningFlagKey(taskID)
+	if timeout == 0 {
+		timeout = DefaultTimeout
+	}
+
 	status := &Status{
 		ID:         taskID,
 		TaskType:   taskType,
@@ -81,63 +89,55 @@ func (s *Task) StartTask(ctx context.Context, keyGroupID uint, taskType, resourc
 		Total:      total,
 		StartedAt:  time.Now(),
 	}
-	statusBytes, err := json.Marshal(status)
-	if err != nil {
-		return nil, fmt.Errorf("failed to serialize new task status: %w", err)
-	}
 
-	if timeout == 0 {
-		timeout = ResultTTL * 24
-	}
-
-	if err := s.store.Set(ctx, lockKey, []byte(taskID), timeout); err != nil {
-		return nil, fmt.Errorf("failed to acquire task resource lock: %w", err)
-	}
-	if err := s.store.Set(ctx, taskKey, statusBytes, timeout); err != nil {
+	if err := s.saveStatus(ctx, taskID, status, timeout); err != nil {
 		_ = s.store.Del(ctx, lockKey)
-		return nil, fmt.Errorf("failed to set new task data in store: %w", err)
+		return nil, fmt.Errorf("failed to save task status: %w", err)
 	}
 
+	runningFlagKey := s.getIsRunningFlagKey(taskID)
 	if err := s.store.Set(ctx, runningFlagKey, []byte("1"), timeout); err != nil {
 		_ = s.store.Del(ctx, lockKey)
-		_ = s.store.Del(ctx, taskKey)
-		return nil, fmt.Errorf("failed to set task running flag: %w", err)
+		_ = s.store.Del(ctx, s.getTaskDataKey(taskID))
+		return nil, fmt.Errorf("failed to set running flag: %w", err)
 	}
+
 	return status, nil
 }
 
 func (s *Task) EndTaskByID(ctx context.Context, taskID, resourceID string, resultData any, taskErr error) {
 	lockKey := s.getResourceLockKey(resourceID)
-	defer func() {
-		if err := s.store.Del(ctx, lockKey); err != nil {
-			s.logger.WithError(err).Warnf("Failed to release resource lock '%s' for task %s.", lockKey, taskID)
-		}
-	}()
 	runningFlagKey := s.getIsRunningFlagKey(taskID)
-	_ = s.store.Del(ctx, runningFlagKey)
+
+	defer func() {
+		_ = s.store.Del(ctx, lockKey)
+		_ = s.store.Del(ctx, runningFlagKey)
+	}()
 
 	status, err := s.GetStatus(ctx, taskID)
 	if err != nil {
-		s.logger.WithError(err).Errorf("Could not get task status for task ID %s during EndTask. Lock has been released, but task data may be stale.", taskID)
+		s.logger.WithError(err).Errorf("Failed to get task status for %s during EndTask", taskID)
 		return
 	}
+
 	if !status.IsRunning {
-		s.logger.Warnf("EndTaskByID called for an already finished task: %s", taskID)
+		s.logger.Warnf("EndTaskByID called for already finished task: %s", taskID)
 		return
 	}
+
 	now := time.Now()
 	status.IsRunning = false
 	status.FinishedAt = &now
 	status.DurationSeconds = now.Sub(status.StartedAt).Seconds()
+
 	if taskErr != nil {
 		status.Error = taskErr.Error()
 	} else {
 		status.Result = resultData
 	}
-	updatedTaskBytes, _ := json.Marshal(status)
-	taskKey := s.getTaskDataKey(taskID)
-	if err := s.store.Set(ctx, taskKey, updatedTaskBytes, ResultTTL); err != nil {
-		s.logger.WithError(err).Errorf("Failed to save final status for task %s.", taskID)
+
+	if err := s.saveStatus(ctx, taskID, status, ResultTTL); err != nil {
+		s.logger.WithError(err).Errorf("Failed to save final status for task %s", taskID)
 	}
 }
 
@@ -148,43 +148,42 @@ func (s *Task) GetStatus(ctx context.Context, taskID string) (*Status, error) {
 		if errors.Is(err, store.ErrNotFound) {
 			return nil, errors.New("task not found")
 		}
-		return nil, fmt.Errorf("failed to get task status from store: %w", err)
+		return nil, fmt.Errorf("failed to get task status: %w", err)
 	}
 
 	var status Status
 	if err := json.Unmarshal(statusBytes, &status); err != nil {
-		return nil, fmt.Errorf("corrupted task data in store for ID %s", taskID)
+		return nil, fmt.Errorf("corrupted task data for ID %s: %w", taskID, err)
 	}
+
 	if !status.IsRunning && status.FinishedAt != nil {
 		status.DurationSeconds = status.FinishedAt.Sub(status.StartedAt).Seconds()
 	}
+
 	return &status, nil
 }
 
 func (s *Task) updateTask(ctx context.Context, taskID string, updater func(status *Status)) error {
 	runningFlagKey := s.getIsRunningFlagKey(taskID)
 	if _, err := s.store.Get(ctx, runningFlagKey); err != nil {
-		return nil
+		if errors.Is(err, store.ErrNotFound) {
+			return errors.New("task is not running")
+		}
+		return fmt.Errorf("failed to check running flag: %w", err)
 	}
+
 	status, err := s.GetStatus(ctx, taskID)
 	if err != nil {
-		s.logger.WithError(err).Warnf("Failed to get task status for update on task %s. Update will not be saved.", taskID)
-		return nil
+		return fmt.Errorf("failed to get task status: %w", err)
 	}
+
 	if !status.IsRunning {
-		return nil
+		return errors.New("task is not running")
 	}
+
 	updater(status)
-	statusBytes, marshalErr := json.Marshal(status)
-	if marshalErr != nil {
-		s.logger.WithError(marshalErr).Errorf("Failed to serialize status for update on task %s.", taskID)
-		return nil
-	}
-	taskKey := s.getTaskDataKey(taskID)
-	if err := s.store.Set(ctx, taskKey, statusBytes, ResultTTL*24); err != nil {
-		s.logger.WithError(err).Warnf("Failed to save update for task %s.", taskID)
-	}
-	return nil
+
+	return s.saveStatus(ctx, taskID, status, DefaultTimeout)
 }
 
 func (s *Task) UpdateProgressByID(ctx context.Context, taskID string, processed int) error {
@@ -198,3 +197,17 @@ func (s *Task) UpdateTotalByID(ctx context.Context, taskID string, total int) er
 		status.Total = total
 	})
 }
+
+func (s *Task) saveStatus(ctx context.Context, taskID string, status *Status, ttl time.Duration) error {
+	statusBytes, err := json.Marshal(status)
+	if err != nil {
+		return fmt.Errorf("failed to serialize status: %w", err)
+	}
+
+	taskKey := s.getTaskDataKey(taskID)
+	if err := s.store.Set(ctx, taskKey, statusBytes, ttl); err != nil {
+		return fmt.Errorf("failed to save status: %w", err)
+	}
+
+	return nil
+}