neutrino + query: Priority index implementation

- The workmanager is made to check if a request has a priority
index before assigning a query index.

- Job index type is also changed to float64 in this commit for
flexibility.

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

Author: Chinwendu20

blockmanager.go

@@ -820,7 +820,7 @@ type checkpointedCFHeadersQuery struct {
 type encodedQuery struct {
 	message       wire.Message
 	encoding      wire.MessageEncoding
-	priorityIndex uint64
+	priorityIndex float64
 }
 
 // Message returns the wire.Message of encodedQuery's struct.
@@ -830,7 +830,7 @@ func (e *encodedQuery) Message() wire.Message {
 
 // PriorityIndex returns the specified priority the caller wants
 // the request to take.
-func (e *encodedQuery) PriorityIndex() uint64 {
+func (e *encodedQuery) PriorityIndex() float64 {
 	return e.priorityIndex
 }
 

query/interface.go

@@ -150,7 +150,7 @@ type Request struct {
 
 type ReqMessage interface {
 	Message() wire.Message
-	PriorityIndex() uint64
+	PriorityIndex() float64
 }
 
 // WorkManager defines an API for a manager that dispatches queries to bitcoin

query/worker.go

@@ -23,7 +23,7 @@ var (
 // addition to some information about the query.
 type queryJob struct {
 	tries      uint8
-	index      uint64
+	index      float64
 	timeout    time.Duration
 	cancelChan <-chan struct{}
 	*Request
@@ -36,7 +36,7 @@ var _ Task = (*queryJob)(nil)
 // Index returns the queryJob's index within the work queue.
 //
 // NOTE: Part of the Task interface.
-func (q *queryJob) Index() uint64 {
+func (q *queryJob) Index() float64 {
 	return q.index
 }
 

query/worker_test.go

@@ -12,14 +12,14 @@ import (
 type mockQueryEncoded struct {
 	message  *wire.MsgGetData
 	encoding wire.MessageEncoding
-	index    uint64
+	index    float64
 }
 
 func (m *mockQueryEncoded) Message() wire.Message {
 	return m.message
 }
 
-func (m *mockQueryEncoded) PriorityIndex() uint64 {
+func (m *mockQueryEncoded) PriorityIndex() float64 {
 	return m.index
 }
 

query/workmanager.go

@@ -199,8 +199,8 @@ func (w *peerWorkManager) workDispatcher() {
 	// We set up a counter that we'll increase with each incoming query,
 	// and will serve as the priority of each. In addition we map each
 	// query to the batch they are part of.
-	queryIndex := uint64(0)
-	currentQueries := make(map[uint64]uint64)
+	queryIndex := float64(0)
+	currentQueries := make(map[float64]uint64)
 
 	workers := make(map[string]*activeWorker)
 
@@ -435,14 +435,25 @@ Loop:
 				"work queue", batchIndex, len(batch.requests))
 
 			for _, q := range batch.requests {
+				idx := queryIndex
+
+				// If priority index is set, use that index.
+				if q.Req.PriorityIndex() != 0 {
+					idx = q.Req.PriorityIndex()
+				}
 				heap.Push(work, &queryJob{
-					index:      queryIndex,
+					index:      idx,
 					timeout:    minQueryTimeout,
 					cancelChan: batch.options.cancelChan,
 					Request:    q,
 				})
-				currentQueries[queryIndex] = batchIndex
-				queryIndex++
+				currentQueries[idx] = batchIndex
+
+				// Only increment queryIndex if it was
+				// assigned to this job.
+				if q.Req.PriorityIndex() == 0 {
+					queryIndex++
+				}
 			}
 
 			currentBatches[batchIndex] = &batchProgress{

query/workmanager_test.go

@@ -131,7 +131,10 @@ func TestWorkManagerWorkDispatcherSingleWorker(t *testing.T) {
 	// Schedule a batch of queries.
 	var queries []*Request
 	for i := 0; i < numQueries; i++ {
-		q := &Request{}
+		q := &Request{
+			Req: &mockQueryEncoded{},
+		}
+
 		queries = append(queries, q)
 	}
 
@@ -141,7 +144,7 @@ func TestWorkManagerWorkDispatcherSingleWorker(t *testing.T) {
 
 	// Each query should be sent on the nextJob queue, in the order they
 	// had in their batch.
-	for i := uint64(0); i < numQueries; i++ {
+	for i := float64(0); i < numQueries; i++ {
 		var job *queryJob
 		select {
 		case job = <-wk.nextJob:
@@ -199,7 +202,9 @@ func TestWorkManagerWorkDispatcherFailures(t *testing.T) {
 	var scheduledJobs [numQueries]chan sched
 	var queries [numQueries]*Request
 	for i := 0; i < numQueries; i++ {
-		q := &Request{}
+		q := &Request{
+			Req: &mockQueryEncoded{},
+		}
 		queries[i] = q
 		scheduledJobs[i] = make(chan sched)
 	}
@@ -212,7 +217,7 @@ func TestWorkManagerWorkDispatcherFailures(t *testing.T) {
 		go func() {
 			for {
 				job := <-wk.nextJob
-				scheduledJobs[job.index] <- sched{
+				scheduledJobs[int(job.index)] <- sched{
 					wk:  wk,
 					job: job,
 				}
@@ -224,11 +229,11 @@ func TestWorkManagerWorkDispatcherFailures(t *testing.T) {
 	errChan := wm.Query(queries[:])
 
 	var jobs [numQueries]sched
-	for i := uint64(0); i < numQueries; i++ {
+	for i := 0; i < numQueries; i++ {
 		var s sched
 		select {
 		case s = <-scheduledJobs[i]:
-			if s.job.index != i {
+			if s.job.index != float64(i) {
 				t.Fatalf("wrong index")
 			}
 
@@ -238,7 +243,7 @@ func TestWorkManagerWorkDispatcherFailures(t *testing.T) {
 			t.Fatalf("next job not received")
 		}
 
-		jobs[s.job.index] = s
+		jobs[int(s.job.index)] = s
 	}
 
 	// Go backwards, and fail half of them.
@@ -262,10 +267,10 @@ func TestWorkManagerWorkDispatcherFailures(t *testing.T) {
 
 	// Finally, make sure the failed jobs are being retried, in the same
 	// order as they were originally scheduled.
-	for i := uint64(0); i < numQueries; i += 2 {
+	for i := float64(0); i < numQueries; i += 2 {
 		var s sched
 		select {
-		case s = <-scheduledJobs[i]:
+		case s = <-scheduledJobs[int(i)]:
 			if s.job.index != i {
 				t.Fatalf("wrong index")
 			}
@@ -309,7 +314,9 @@ func TestWorkManagerCancelBatch(t *testing.T) {
 	// Schedule a batch of queries.
 	var queries []*Request
 	for i := 0; i < numQueries; i++ {
-		q := &Request{}
+		q := &Request{
+			Req: &mockQueryEncoded{},
+		}
 		queries = append(queries, q)
 	}
 
@@ -399,7 +406,9 @@ func TestWorkManagerWorkRankingScheduling(t *testing.T) {
 	// Schedule a batch of queries.
 	var queries []*Request
 	for i := 0; i < numQueries; i++ {
-		q := &Request{}
+		q := &Request{
+			Req: &mockQueryEncoded{},
+		}
 		queries = append(queries, q)
 	}
 
@@ -411,7 +420,7 @@ func TestWorkManagerWorkRankingScheduling(t *testing.T) {
 	for i := 0; i < numQueries; i++ {
 		select {
 		case job := <-workers[i].nextJob:
-			if job.index != uint64(i) {
+			if job.index != float64(i) {
 				t.Fatalf("unexpected job")
 			}
 			jobs = append(jobs, job)
@@ -462,7 +471,9 @@ func TestWorkManagerWorkRankingScheduling(t *testing.T) {
 	// Send a new set of queries.
 	queries = nil
 	for i := 0; i < numQueries; i++ {
-		q := &Request{}
+		q := &Request{
+			Req: &mockQueryEncoded{},
+		}
 		queries = append(queries, q)
 	}
 	_ = wm.Query(queries)
@@ -476,3 +487,121 @@ func TestWorkManagerWorkRankingScheduling(t *testing.T) {
 		}
 	}
 }
+
+// TestWorkManagerSchedulePriorityIndex tests that the workmanager acknowledges
+// priority index.
+func TestWorkManagerSchedulePriorityIndex(t *testing.T) {
+	const numQueries = 3
+
+	// Start work manager with as many workers as queries. This is not very
+	// realistic, but makes the work manager able to schedule all queries
+	// concurrently.
+	wm, workers := startWorkManager(t, numQueries)
+
+	// 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 [5]chan sched
+	var queries [numQueries]*Request
+	for i := 0; i < numQueries; i++ {
+		var q *Request
+		idx := i
+		if i == 0 {
+			q = &Request{
+				Req: &mockQueryEncoded{},
+			}
+		} else {
+			// Assign priority index.
+			idx = i + 2
+			q = &Request{
+				Req: &mockQueryEncoded{
+					index: float64(idx),
+				},
+			}
+		}
+		queries[i] = q
+		scheduledJobs[idx] = 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 jobs [numQueries]sched
+	for i := uint64(0); i < numQueries; i++ {
+		var expectedIndex float64
+
+		if i == 0 {
+			expectedIndex = float64(0)
+		} else {
+			expectedIndex = float64(i + 2)
+		}
+		var s sched
+		select {
+		case s = <-scheduledJobs[int(expectedIndex)]:
+
+			if s.job.index != expectedIndex {
+				t.Fatalf("wrong index: Got %v but expected %v", s.job.index,
+					expectedIndex)
+			}
+		case <-errChan:
+			t.Fatalf("did not expect an errChan")
+		case <-time.After(time.Second):
+			t.Fatalf("next job not received")
+		}
+
+		jobs[i] = s
+	}
+
+	// Go backwards send results for job.
+	for i := numQueries - 1; i >= 0; i-- {
+		select {
+		case jobs[i].wk.results <- &jobResult{
+			job: jobs[i].job,
+		}:
+		case <-errChan:
+			t.Fatalf("did not expect on errChan")
+		case <-time.After(time.Second):
+			t.Fatalf("result not handled")
+		}
+	}
+
+	// Finally, make sure no jobs are retried.
+	for i := uint64(0); i < numQueries; i++ {
+		select {
+		case <-scheduledJobs[i]:
+			t.Fatalf("did not expect a retried job")
+		case <-time.After(time.Second):
+		}
+	}
+
+	// The query should ultimately succeed.
+	select {
+	case err := <-errChan:
+		if err != nil {
+			t.Fatalf("got error: %v", err)
+		}
+	case <-time.After(time.Second):
+		t.Fatalf("nothing received on errChan")
+	}
+}

query/workqueue.go

@@ -4,7 +4,7 @@ package query
 // work queue.
 type Task interface {
 	// Index returns this Task's index in the work queue.
-	Index() uint64
+	Index() float64
 }
 
 // workQueue is struct implementing the heap interface, and is used to keep a