From ddd1f21e5d516db63a07999ac81eeb6450504fd5 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=C5=81ukasz=20Magiera?= <magik6k@gmail.com>
Date: Wed, 24 Jun 2020 23:06:56 +0200
Subject: [PATCH] sched: Support external priority

---
 request_queue.go |  4 ++++
 sched.go         | 20 ++++++++++++++++++++
 sched_test.go    | 18 ++++++++++++++++++
 3 files changed, 42 insertions(+)
 create mode 100644 sched_test.go

diff --git a/request_queue.go b/request_queue.go
index e5b3fd234..0d35e9f1d 100644
--- a/request_queue.go
+++ b/request_queue.go
@@ -7,6 +7,10 @@ type requestQueue []*workerRequest
 func (q requestQueue) Len() int { return len(q) }
 
 func (q requestQueue) Less(i, j int) bool {
+	if q[i].priority != q[j].priority {
+		return q[i].priority > q[j].priority
+	}
+
 	if q[i].taskType != q[j].taskType {
 		return q[i].taskType.Less(q[j].taskType)
 	}
diff --git a/sched.go b/sched.go
index 3b3e28d65..8920a1a21 100644
--- a/sched.go
+++ b/sched.go
@@ -16,6 +16,24 @@ import (
 	"github.com/filecoin-project/sector-storage/storiface"
 )
 
+type schedPrioCtxKey int
+
+var SchedPriorityKey schedPrioCtxKey
+var DefaultSchedPriority = 0
+
+func getPriority(ctx context.Context) int {
+	sp := ctx.Value(SchedPriorityKey)
+	if p, ok := sp.(int); ok {
+		return p
+	}
+
+	return DefaultSchedPriority
+}
+
+func WithPriority(ctx context.Context, priority int) context.Context {
+	return context.WithValue(ctx, SchedPriorityKey, priority)
+}
+
 const mib = 1 << 20
 
 type WorkerAction func(ctx context.Context, w Worker) error
@@ -72,6 +90,7 @@ func (sh *scheduler) Schedule(ctx context.Context, sector abi.SectorID, taskType
 	case sh.schedule <- &workerRequest{
 		sector:   sector,
 		taskType: taskType,
+		priority: getPriority(ctx),
 		sel:      sel,
 
 		prepare: prepare,
@@ -99,6 +118,7 @@ func (sh *scheduler) Schedule(ctx context.Context, sector abi.SectorID, taskType
 type workerRequest struct {
 	sector   abi.SectorID
 	taskType sealtasks.TaskType
+	priority int // larger values more important
 	sel      WorkerSelector
 
 	prepare WorkerAction
diff --git a/sched_test.go b/sched_test.go
new file mode 100644
index 000000000..d0d0e7ca9
--- /dev/null
+++ b/sched_test.go
@@ -0,0 +1,18 @@
+package sectorstorage
+
+import (
+	"context"
+	"testing"
+
+	"github.com/stretchr/testify/require"
+)
+
+func TestWithPriority(t *testing.T) {
+	ctx := context.Background()
+
+	require.Equal(t, DefaultSchedPriority, getPriority(ctx))
+
+	ctx = WithPriority(ctx, 2222)
+
+	require.Equal(t, 2222, getPriority(ctx))
+}