From 4c26a9bc18c95edbcf39f8e51824950140eb684b Mon Sep 17 00:00:00 2001
From: Aaditya Sondhi <20070511+aadityasondhi@users.noreply.github.com>
Date: Tue, 12 Dec 2023 21:55:24 +0000
Subject: [PATCH] kv,admission: add missing admission headers

This patch adds missing admission headers in a couple of places:
`txn_interceptor_heartbeater` and `replica_range_lease`. Batch requests
coming from these code paths are expected to bypass AC. Empty admission
headers were accomplishing the same goal but this makes it more clear at
what the intention is.

Informs #112680

Release note: None
---
 pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go | 7 +++++++
 pkg/kv/kvserver/replica_range_lease.go                 | 8 ++++++++
 2 files changed, 15 insertions(+)

diff --git a/pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go b/pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go
index 38de47c9bf5..959aa357de2 100644
--- a/pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go
+++ b/pkg/kv/kvclient/kvcoord/txn_interceptor_heartbeater.go
@@ -525,6 +525,13 @@ func (h *txnHeartbeater) abortTxnAsyncLocked(ctx context.Context) {
 		// concurrent requests from failing to notice the transaction was aborted.
 		Poison: true,
 	})
+	// NB: Setting `Source: kvpb.AdmissionHeader_OTHER` means this request will
+	// bypass AC.
+	ba.AdmissionHeader = kvpb.AdmissionHeader{
+		Priority:   txn.AdmissionPriority,
+		CreateTime: timeutil.Now().UnixNano(),
+		Source:     kvpb.AdmissionHeader_OTHER,
+	}
 
 	const taskName = "txnHeartbeater: aborting txn"
 	log.VEventf(ctx, 2, "async abort for txn: %s", txn)
diff --git a/pkg/kv/kvserver/replica_range_lease.go b/pkg/kv/kvserver/replica_range_lease.go
index 407e6c149c1..f2924ec801b 100644
--- a/pkg/kv/kvserver/replica_range_lease.go
+++ b/pkg/kv/kvserver/replica_range_lease.go
@@ -58,6 +58,7 @@ import (
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/util"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
 	"github.com/cockroachdb/cockroach/pkg/util/envutil"
 	"github.com/cockroachdb/cockroach/pkg/util/growstack"
 	"github.com/cockroachdb/cockroach/pkg/util/hlc"
@@ -571,6 +572,13 @@ func (p *pendingLeaseRequest) requestLease(
 	// lease when the range is unavailable results in, essentially, giving
 	// up on the lease and thus worsening the situation.
 	ba.Add(leaseReq)
+	// NB: Setting `Source: kvpb.AdmissionHeader_OTHER` means this request will
+	// bypass AC.
+	ba.AdmissionHeader = kvpb.AdmissionHeader{
+		Priority:   int32(admissionpb.NormalPri),
+		CreateTime: timeutil.Now().UnixNano(),
+		Source:     kvpb.AdmissionHeader_OTHER,
+	}
 	_, pErr := p.repl.Send(ctx, ba)
 	return pErr.GoError()
 }