query: ErrQueryTimeout does not exit worker feedback loop.

Worker waits after timeout for response but job is scheduled on
another worker so as not to be slowed down by one worker.

Signed-off-by: Maureen Ononiwu <[email protected]>

Author: Chinwendu20

query.go

@@ -430,6 +430,7 @@ type cfiltersQuery struct {
 	headerIndex   map[chainhash.Hash]int
 	targetHash    chainhash.Hash
 	targetFilter  *gcs.Filter
+	mtx           sync.Mutex
 }
 
 // request couples a query message with the handler to be used for the response
@@ -498,6 +499,8 @@ func (q *cfiltersQuery) handleResponse(req, resp wire.Message,
 
 	// If this filter is for a block not in our index, we can ignore it, as
 	// we either already got it, or it is out of our queried range.
+	q.mtx.Lock()
+	defer q.mtx.Unlock()
 	i, ok := q.headerIndex[response.BlockHash]
 	if !ok {
 		return noProgress

query/worker.go

@@ -134,7 +134,7 @@ nexJobLoop:
 			err := job.SendQuery(peer, job.Req)
 
 			// If any error occurs while sending query send a message to the worker's feedback Chan
-			// which would be handled by the "LOOP" below.
+			// which would be handled by the "feedbackLoop" below.
 			if err != nil {
 				select {
 				case results <- &jobResult{
@@ -268,6 +268,16 @@ nexJobLoop:
 			return
 		}
 
+		// If the error is a timeout still wait for the response as we are assured a response as long as there was a
+		// request but reschedule on another worker to quickly fetch a response so as not to be slowed down by this
+		// worker. We either get a response or the peer stalls (i.e. disconnects due to an elongated time without
+		// a response)
+		if jobErr == ErrQueryTimeout {
+			jobErr = nil
+
+			goto feedbackLoop
+		}
+
 		// If the peer disconnected, we can exit immediately.
 		if jobErr == ErrPeerDisconnected {
 			return

query/worker_test.go

@@ -303,8 +303,8 @@ func TestWorkerTimeout(t *testing.T) {
 	// It will immediately attempt to fetch another task.
 	select {
 	case ctx.nextJob <- task:
+		t.Fatalf("worker still in feedback loop picked up job")
 	case <-time.After(1 * time.Second):
-		t.Fatalf("did not pick up job")
 	}
 }
 

query/workmanager.go

@@ -298,7 +298,9 @@ Loop:
 			// Delete the job from the worker's active job, such
 			// that the slot gets opened for more work.
 			r := workers[result.peer.Addr()]
-			r.activeJob = nil
+			if result.err != ErrQueryTimeout {
+				r.activeJob = nil
+			}
 
 			// Get the index of this query's batch, and delete it
 			// from the map of current queries, since we don't have

query/workmanager_test.go

@@ -302,6 +302,106 @@ func TestWorkManagerWorkDispatcherFailures(t *testing.T) {
 	}
 }
 
+// TestWorkManagerErrQueryTimeout tests that the workers that return query
+// timeout are not sent jobs until they return a different error.
+func TestWorkManagerErrQueryTimeout(t *testing.T) {
+	const numQueries = 2
+	const numWorkers = 1
+
+	// Start work manager.
+	wm, workers := startWorkManager(t, numWorkers)
+
+	// When the jobs gets scheduled, keep track of which worker was
+	// assigned the job.
+	type sched struct {
+		wk  *mockWorker
+		job *queryJob
+	}
+
+	// Schedule a batch of queries.
+	var scheduledJobs [numQueries]chan sched
+	var queries [numQueries]*Request
+	for i := 0; i < numQueries; i++ {
+		q := &Request{
+			Req: &mockQueryEncoded{},
+		}
+		queries[i] = q
+		scheduledJobs[i] = make(chan sched)
+	}
+
+	// Fot each worker, spin up a goroutine that will forward the job it
+	// got to our slice of scheduled jobs, such that we can handle them in
+	// order.
+	for i := 0; i < len(workers); i++ {
+		wk := workers[i]
+		go func() {
+			for {
+				job := <-wk.nextJob
+				scheduledJobs[int(job.index)] <- sched{
+					wk:  wk,
+					job: job,
+				}
+			}
+		}()
+	}
+
+	// Send the batch, and Retrieve all jobs immediately.
+	errChan := wm.Query(queries[:])
+
+	var iter int
+	var s sched
+	for i := 0; i < numQueries; i++ {
+		select {
+		case s = <-scheduledJobs[i]:
+			if s.job.index != float64(i) {
+				t.Fatalf("wrong index")
+			}
+			if iter == 1 {
+				t.Fatalf("Expected only one scheduled job")
+			}
+			iter++
+		case <-errChan:
+			t.Fatalf("did not expect on errChan")
+		case <-time.After(time.Second):
+			if iter < 1 {
+				t.Fatalf("next job not received")
+			}
+		}
+	}
+
+	select {
+	case s.wk.results <- &jobResult{
+		job: s.job,
+		err: ErrQueryTimeout,
+	}:
+	case <-errChan:
+		t.Fatalf("did not expect on errChan")
+	case <-time.After(time.Second):
+		t.Fatalf("result not handled")
+	}
+
+	// Finally, make sure the job is not retried as there are no available
+	// peer to retry it.
+
+	select {
+	case <-scheduledJobs[0]:
+		t.Fatalf("did not expect job rescheduled")
+	case <-errChan:
+		t.Fatalf("did not expect on errChan")
+	case <-time.After(time.Second):
+	}
+
+	// There should be no errChan message as query is still incomplete.
+	select {
+	case err := <-errChan:
+		if err != nil {
+			t.Fatalf("got error: %v", err)
+		}
+		t.Fatalf("expected no errChan message")
+	case <-time.After(time.Second):
+	}
+}
+
 // TestWorkManagerCancelBatch checks that we can cancel a batch query midway,
 // and that the jobs it contains are canceled.
 func TestWorkManagerCancelBatch(t *testing.T) {