Fix potential task requeueing failure (#76)

Author: man90es

handlers/getStatus.go

@@ -12,7 +12,7 @@ import (
 )
 
 var initTime = time.Now()
-var version = "1.19.0"
+var version = "1.19.1"
 
 func getStatus(w http.ResponseWriter, r *http.Request) {
 	json.NewEncoder(w).Encode(map[string]interface{}{

scraper/handleTaskError.go

@@ -19,6 +19,24 @@ func handleTaskError(r *colly.Request, blocked bool, err error) {
 	taskClient := r.Ctx.Get(metadataTaskClient)
 	taskHash := r.Ctx.Get(metadataTaskHash)
 
+	taskQueue.ConfirmTaskCompletion(taskClient, taskHash)
+	if taskRetries < viper.GetInt("taskretries") {
+		taskQueue.AddTask(
+			taskClient,
+			taskHash,
+			r.URL.String(),
+			viper.GetBool("taskretryfront"),
+			map[string]string{
+				metadataTaskAddedAt: r.Ctx.Get(metadataTaskAddedAt),
+				metadataTaskClient:  taskClient,
+				metadataTaskHash:    taskHash,
+				metadataTaskRegion:  r.Ctx.Get(metadataTaskRegion),
+				metadataTaskRetries: strconv.Itoa(taskRetries + 1),
+				metadataTaskType:    r.Ctx.Get(metadataTaskType),
+			},
+		)
+	}
+
 	if blocked {
 		logger.Error(fmt.Sprintf("Hit Imperva while loading %v, retries: %v", r.URL, taskRetries))
 	} else if strings.Contains(err.Error(), "http2: Transport received GOAWAY from server ErrCode:INTERNAL_ERROR") {
@@ -40,23 +58,4 @@ func handleTaskError(r *colly.Request, blocked bool, err error) {
 	} else {
 		taskQueue.Pause(time.Duration(60-time.Now().Second()) * time.Second)
 	}
-
-	taskQueue.ConfirmTaskCompletion(taskClient, taskHash)
-
-	if taskRetries < viper.GetInt("taskretries") {
-		taskQueue.AddTask(
-			taskClient,
-			taskHash,
-			r.URL.String(),
-			viper.GetBool("taskretryfront"),
-			map[string]string{
-				metadataTaskAddedAt: r.Ctx.Get(metadataTaskAddedAt),
-				metadataTaskClient:  taskClient,
-				metadataTaskHash:    taskHash,
-				metadataTaskRegion:  r.Ctx.Get(metadataTaskRegion),
-				metadataTaskRetries: strconv.Itoa(taskRetries + 1),
-				metadataTaskType:    r.Ctx.Get(metadataTaskType),
-			},
-		)
-	}
 }

scraper/taskQueue.go

@@ -28,6 +28,15 @@ var enqueueScript = redis.NewScript(`
 	return 1
 `)
 
+var confirmCompletionScript = redis.NewScript(`
+	local clientsKey, hashesKey = KEYS[1], KEYS[2]
+	local taskClient, hash = ARGV[1], ARGV[2]
+
+	redis.call("HINCRBY", clientsKey, taskClient, -1)
+	redis.call("SREM", hashesKey, hash)
+	return 1
+`)
+
 type Task struct {
 	Hash       string
 	Metadata   map[string]string
@@ -198,8 +207,5 @@ func (q *TaskQueue) CountQueuedTasksForClient(taskClient string) (count int) {
 }
 
 func (q *TaskQueue) ConfirmTaskCompletion(taskClient string, hash string) {
-	pipe := q.rdb.Pipeline()
-	pipe.HIncrBy(q.ctx, q.clientsKey, taskClient, -1)
-	pipe.SRem(q.ctx, q.hashesKey, hash)
-	_, _ = pipe.Exec(q.ctx)
+	_, _ = confirmCompletionScript.Run(q.ctx, q.rdb, []string{q.clientsKey, q.hashesKey}, taskClient, hash).Result()
 }