From 3f1747d212b9471b96e37c965380893f231e5647 Mon Sep 17 00:00:00 2001
From: Aaron Beitch <aaronb@arista.com>
Date: Wed, 14 Aug 2024 11:42:51 -0700
Subject: [PATCH] hrpc, region: Allow setting priority on Scan and Get requests

Provide an option for setting the priority on a Scan or Get request.
The priority is set on the RequestHeader, so must be set by the region
code when serializing the request.
---
 hrpc/call.go       |  1 +
 hrpc/query.go      | 20 ++++++++++++++++++++
 hrpc/query_test.go | 46 ++++++++++++++++++++++++++++++++++++++++++++++
 region/client.go   |  5 ++++-
 4 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/hrpc/call.go b/hrpc/call.go
index a15e6e4c..7847b914 100644
--- a/hrpc/call.go
+++ b/hrpc/call.go
@@ -101,6 +101,7 @@ type hasQueryOptions interface {
 	setResultOffset(offset uint32)
 	setCacheBlocks(cacheBlocks bool)
 	setConsistency(consistency ConsistencyType)
+	setPriority(priority uint32)
 }
 
 // RPCResult is struct that will contain both the resulting message from an RPC
diff --git a/hrpc/query.go b/hrpc/query.go
index c652819d..d2798907 100644
--- a/hrpc/query.go
+++ b/hrpc/query.go
@@ -23,6 +23,7 @@ type baseQuery struct {
 	maxVersions   uint32
 	storeLimit    uint32
 	storeOffset   uint32
+	priority      uint32
 	cacheBlocks   bool
 	consistency   ConsistencyType
 }
@@ -97,6 +98,15 @@ func (bq *baseQuery) setCacheBlocks(cacheBlocks bool) {
 func (bq *baseQuery) setConsistency(consistency ConsistencyType) {
 	bq.consistency = consistency
 }
+func (bq *baseQuery) setPriority(priority uint32) {
+	bq.priority = priority
+}
+func (bq *baseQuery) Priority() *uint32 {
+	if bq.priority == 0 {
+		return nil
+	}
+	return &bq.priority
+}
 
 // Families option adds families constraint to a Scan or Get request.
 func Families(f map[string][]string) func(Call) error {
@@ -218,3 +228,13 @@ func Consistency(consistency ConsistencyType) func(Call) error {
 		return errors.New("'Consistency' option can only be used with Get or Scan requests")
 	}
 }
+
+func Priority(priority uint32) func(Call) error {
+	return func(hc Call) error {
+		if c, ok := hc.(hasQueryOptions); ok {
+			c.setPriority(priority)
+			return nil
+		}
+		return errors.New("'Priority' option can only be used with Get or Scan requests")
+	}
+}
diff --git a/hrpc/query_test.go b/hrpc/query_test.go
index a1daf04b..884eb13a 100644
--- a/hrpc/query_test.go
+++ b/hrpc/query_test.go
@@ -208,3 +208,49 @@ func TestCacheBlocks(t *testing.T) {
 		t.Error(err)
 	}
 }
+
+func TestPriority(t *testing.T) {
+	get, err := NewGet(nil, nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := get.Priority(); got != nil {
+		t.Errorf("expected nil, got %v", got)
+	}
+	get, err = NewGet(nil, nil, nil, Priority(0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := get.Priority(); got != nil {
+		t.Errorf("expected nil, got %v", got)
+	}
+	get, err = NewGet(nil, nil, nil, Priority(5))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := get.Priority(); *got != 5 {
+		t.Errorf("expected priority 5, got %v", got)
+	}
+
+	scan, err := NewScan(nil, nil)
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := scan.Priority(); got != nil {
+		t.Errorf("expected nil, got %v", got)
+	}
+	scan, err = NewScan(nil, nil, Priority(0))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := scan.Priority(); got != nil {
+		t.Errorf("expected nil, got %v", got)
+	}
+	scan, err = NewScan(nil, nil, Priority(5))
+	if err != nil {
+		t.Fatal(err)
+	}
+	if got := scan.Priority(); *got != 5 {
+		t.Errorf("expected priority 5, got %v", got)
+	}
+}
diff --git a/region/client.go b/region/client.go
index bb6f1e36..c50ae345 100644
--- a/region/client.go
+++ b/region/client.go
@@ -683,10 +683,13 @@ var pbTrue = proto.Bool(true)
 func marshalProto(rpc hrpc.Call, callID uint32, request proto.Message,
 	cellblocksLen uint32) ([]byte, error) {
 	header := getHeader()
+	defer returnHeader(header)
 	header.MethodName = proto.String(rpc.Name())
 	header.RequestParam = pbTrue
 	header.CallId = &callID
-	defer returnHeader(header)
+	if p, ok := rpc.(interface{ Priority() *uint32 }); ok {
+		header.Priority = p.Priority()
+	}
 
 	if cellblocksLen > 0 {
 		header.CellBlockMeta = &pb.CellBlockMeta{