admission: support non-blocking {Store,}WorkQueue.Admit()

Part of cockroachdb#95563. For end-to-end flow control of replicated writes, we
want to enable below-raft admission control through the following API on
kvadmission.Controller:

  // AdmitRaftEntry informs admission control of a raft log entry being
  // written to storage (for the given tenant, the specific range, and
  // on the named store).
  AdmitRaftEntry(
    context.Context, roachpb.TenantID,
    roachpb.StoreID, roachpb.RangeID, raftpb.Entry.
  )

This serves as the integration point for log entries received below
raft right as they're being written to stable storage. It's a
non-blocking interface since we're below-raft and in the raft.Ready()
loop. What it effectively does is enqueues a "virtual" work item in the
underlying StoreWorkQueue mediating all store IO. This virtual work item
is what later gets dequeued once the IO granter informs the work queue
of newly available IO tokens. When enqueueing the virtual work item, we
still update the StoreWorkQueue's physically-accounted-for bytes since
the actual write is not deferred, and timely statistic updates improves
accuracy for the underlying linear models (that map between
accounted-for writes and observed L0 growth, using it to inform IO token
generation rates).

For each of the arguments above:
- The roachpb.TenantID is plumbed to find the right tenant heap to queue
  it under (for inter-tenant isolation).
- The roachpb.StoreID to find the right store work queue on multi-store
  nodes. We'll also use the StoreID when informing the origin node of
  log entries being admitted[^1].
- We pass in the roachpb.RangeID on behalf of which work is being
  admitted. This, along side the raftpb.Entry.{Term,Index} for the
  replicated write uniquely identifies where the write is to end up.
  We use these identifiers to:
  - Return flow tokens on the origin node[^1][^2].
  - In WorkQueue ordering -- for replicated writes below-raft, we ignore
    CreateTime/epoch-LIFO, and instead sort by priority and within a
    priority, sort by log position.
- For standard work queue ordering, our work item needs to include the
  CreateTime and AdmissionPriority, details that are passed down using
  dedicated raft log entry encodings[^3][^4] as part of the raftpb.Entry
  parameter above.
  - Since the raftpb.Entry encodes within it its origin node[^4], it
    will be used post-admission to dispatch flow tokens to the right
    node. This integration is left to future PRs.

We use the above to populate the following fields on a per-(replicated
write)work basis:

    // ReplicatedWorkInfo groups everything needed to admit replicated
    // writes, done so asynchronously below-raft as part of replication
    // admission control.
    type ReplicatedWorkInfo struct {
      RangeID roachpb.RangeID
      Origin roachpb.NodeID
      LogPosition LogPosition
      Ingested bool
    }

Since admission is happening below-raft where the size of the write is
known, we no longer need per-work estimates for upfront IO token
deductions. Since admission is asynchronous, we also don't use
the AdmittedWorkDone interface which was to make token adjustments
(without blocking) given the upfront estimates. We still want to
intercept the exact point when some write work gets admitted in order to
inform the origin node so it can release flow tokens. We do so through
the following interface satisfied by the StoreWorkQueue:

  // onAdmittedReplicatedWork is used to intercept the
  // point-of-admission for replicated writes.
  type onAdmittedReplicatedWork interface {
    admittedReplicatedWork(
      tenantID roachpb.TenantID,
      pri admissionpb.WorkPriority,
      rwi ReplicatedWorkInfo,
      requestedTokens int64,
    )
  }

[^1]: See kvflowcontrolpb.AdmittedRaftLogEntries introduced in cockroachdb#95637.
[^2]: See kvflowcontrol.Handle.{ReturnTokensUpto,DeductTokensFor}
      introduced in cockroachdb#95637. Token deductions and returns are tied to
      raft log positions.
[^3]: See raftlog.EntryEncoding{Standard,Sideloaded}WithAC introduced in
      cockroachdb#95748.
[^4]: See kvflowcontrolpb.RaftAdmissionMeta introduced in cockroachdb#95637.
        message RaftAdmissionMeta {
          int32 admission_priority = ...;
          int64 admission_create_time = ...;
          int32 admission_origin_node = ...;
        }

Release note: None

pkg/kv/kvserver/kvadmission/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/kv/kvpb",
+        "//pkg/kv/kvserver/raftlog",
         "//pkg/roachpb",
         "//pkg/settings",
         "//pkg/settings/cluster",

pkg/kv/kvserver/kvadmission/kvadmission.go

@@ -19,6 +19,7 @@ import (
 	"time"
 
 	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/raftlog"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
@@ -120,7 +121,7 @@ type Controller interface {
 	FollowerStoreWriteBytes(roachpb.StoreID, FollowerStoreWriteBytes)
 	// AdmitRaftEntry informs admission control of a raft log entry being
 	// written to storage.
-	AdmitRaftEntry(roachpb.TenantID, roachpb.StoreID, roachpb.RangeID, raftpb.Entry)
+	AdmitRaftEntry(context.Context, roachpb.TenantID, roachpb.StoreID, roachpb.RangeID, raftpb.Entry)
 }
 
 // TenantWeightProvider can be periodically asked to provide the tenant
@@ -244,7 +245,7 @@ func (n *controllerImpl) AdmitKVWork(
 			if err != nil {
 				return Handle{}, err
 			}
-			admissionEnabled = storeWorkHandle.AdmissionEnabled()
+			admissionEnabled = storeWorkHandle.UseAdmittedWorkDone()
 			if admissionEnabled {
 				defer func() {
 					if retErr != nil {
@@ -409,9 +410,60 @@ func (n *controllerImpl) FollowerStoreWriteBytes(
 
 // AdmitRaftEntry implements the Controller interface.
 func (n *controllerImpl) AdmitRaftEntry(
-	roachpb.TenantID, roachpb.StoreID, roachpb.RangeID, raftpb.Entry,
+	ctx context.Context,
+	tenantID roachpb.TenantID,
+	storeID roachpb.StoreID,
+	rangeID roachpb.RangeID,
+	entry raftpb.Entry,
 ) {
-	panic("unimplemented")
+	typ, err := raftlog.EncodingOf(entry)
+	if err != nil {
+		log.Errorf(ctx, "unable to determine raft command encoding: %v", err)
+		return
+	}
+	if !typ.UsesAdmissionControl() {
+		return // nothing to do
+	}
+	meta, err := raftlog.DecodeRaftAdmissionMeta(entry.Data)
+	if err != nil {
+		log.Errorf(ctx, "unable to decode raft command admission data: %v", err)
+		return
+	}
+
+	storeAdmissionQ := n.storeGrantCoords.TryGetQueueForStore(int32(storeID))
+	if storeAdmissionQ == nil {
+		log.Errorf(ctx, "unable to find queue for store: %s", storeID)
+		return // nothing to do
+	}
+
+	wi := admission.WorkInfo{
+		TenantID:        tenantID,
+		Priority:        admissionpb.WorkPriority(meta.AdmissionPriority),
+		CreateTime:      meta.AdmissionCreateTime,
+		BypassAdmission: false,
+		RequestedCount:  int64(entry.Size()),
+	}
+	wi.ReplicatedWorkInfo = admission.ReplicatedWorkInfo{
+		Enabled: true,
+		RangeID: rangeID,
+		Origin:  meta.AdmissionOriginNode,
+		LogPosition: admission.LogPosition{
+			Term:  entry.Term,
+			Index: entry.Index,
+		},
+		Ingested: typ.IsSideloaded(),
+	}
+
+	handle, err := storeAdmissionQ.Admit(ctx, admission.StoreWriteWorkInfo{
+		WorkInfo: wi,
+	})
+	if err != nil {
+		log.Errorf(ctx, "error while admitting to store admission queue: %v", err)
+		return
+	}
+	if handle.UseAdmittedWorkDone() {
+		log.Fatalf(ctx, "unexpected handle.UseAdmittedWorkDone")
+	}
 }
 
 // FollowerStoreWriteBytes captures stats about writes done to a store by a

pkg/kv/kvserver/kvflowcontrol/doc.go

@@ -327,6 +327,76 @@ package kvflowcontrol
 //   it can transition into the mode described in I3a where we deduct/block for
 //   flow tokens for subsequent quorum writes.
 //
+// I12. How does this interact with epoch-LIFO? Or CreateTime ordering
+//      generally?
+// - Background: Epoch-LIFO tries to reduce lower percentile admission queueing
+//   delay (at the expense of higher percentile delay) by switching between the
+//   standard CreateTime-based FIFO ordering to LIFO-within-an-epoch. Work
+//   submitted by a transaction with CreateTime T is slotted into the epoch
+//   number with time interval [E, E+100ms) where T is contained in this
+//   interval. The work will start executing after the epoch is closed, at
+//   ~E+100ms. This increases transaction latency by at most ~100ms. When the
+//   epoch closes, all nodes start executing the transactions in that epoch in
+//   LIFO order, and will implicitly have the same set of competing
+//   transactions[^10], a set that stays unchanged until the next epoch closes.
+//   And by the time the next epoch closes, and the current epoch's transactions
+//   are deprioritized, 100ms will have elapsed, which is selected to be long
+//   enough for most of these now-admitted to have finished their work.
+//   - We switch to LIFO-within-an-epoch once we start observing that the
+//     maximum queuing delay for work within a <tenant,priority> starts
+//     exceeding the ~100ms we'd add with epoch-LIFO.
+// - For below-raft work queue ordering, we ignore CreateTime when ordering work
+//   within the same range. Within a given <tenant,priority,range>, admission
+//   takes place in raft log order (i.e. entries with lower terms get admitted
+//   first, or lower indexes within the same term).
+//   - NB: Regarding "admission takes place in raft log order", we can implement
+//     this differently. We introduced log-position based ordering to simplify
+//     the implementation of token returns where we release tokens by specifying
+//     the log position up-to-which we want to release held tokens[^11]. But
+//     with additional tracking in the below-raft work queues, if we know that
+//     work W2 with log position L2 got admitted, and corresponded to F flow
+//     token deductions at the origin, and we also know that work W1 with log
+//     position L1 is currently queued, also corresponding to F flow token
+//     deductions at the origin, we could inform the origin node to return flow
+//     tokens up to L1 and still get what we want -- a return of F flow tokens
+//     when each work gets admitted.
+// - To operate within cluster-wide FIFO ordering, we order by CreateTime when
+//   comparing work across different ranges. Writes for a single range, as
+//   observed by a given store below-raft (follower or otherwise) travel along
+//   a single stream. Consider the case where a single store S3 receives
+//   replication traffic for two ranges R1 and R2, originating from two separate
+//   nodes N1 and N2. If N1 is issuing writes with strictly older CreateTimes,
+//   when returning flow tokens we should prefer N1.
+//   - What about CreateTime ordering within a <tenant,priority,range>? Flow
+//     tokens deductions aren't tied to create times -- they're tied to work
+//     classes on the sender. So we still want priority-based ordering to
+//     release regular flow tokens before elastic ones, but releasing flow
+//     tokens for work with lower CreateTimes does not actually promote doing
+//     older work. Below-raft admission is all asynchronous. To get back to
+//     CreateTime ordering, we'd need to do it above-raft, by introducing a
+//     WorkQueue-like structure for requests waiting for flow tokens.
+//   - We could do the same for epoch-LIFO, but implementing it above-raft in
+//     WorkQueue-like structures where requests wait for flow tokens.
+//     Cluster-wide we'd allow work within a given epoch to start issuing
+//     replication traffic.
+//   - How does this interact with the CreateTime ordering below-raft when
+//     comparing work across different ranges?
+//     - When there's no epoch-LIFO happening, we have cluster-wide FIFO
+//       ordering within a <tenant,priority> pair as described above.
+//     - When epoch-LIFO is happening across all senders issuing replication
+//       traffic to a given receiver store, we're seeing work within the same
+//       epoch, but we'll be returning flow tokens to nodes issuing work with
+//       older CreateTimes. So defeating the LIFO in epoch-LIFO.
+//     - When epoch-LIFO is happening only on a subset of the senders issuing
+//       replication traffic, we'll again be returning flow tokens to nodes
+//       issuing work with older CreateTimes. Which is undefined.
+//     - Is it strange that different nodes can admit work from "different
+//       epochs"? What are we to do below-raft, when deciding where to return flow
+//       tokens back to, since it's all for the same <tenant,priority>? Maybe we
+//       need to pass down whether the work was admitted at the sender with
+//       LIFO/FIFO, and return flow tokens in {LIFO,FIFO} order across all nodes
+//       that issued {LIFO,FIFO} work?
+//
 // ---
 //
 // [^1]: kvserverpb.RaftMessageRequest is the unit of what's sent
@@ -373,6 +443,9 @@ package kvflowcontrol
 //         machine application get significantly behind due to local scheduling
 //         reasons by using the same goroutine to do both async raft log writes
 //         and state machine application.
+// [^10]: This relies on partially synchronized clocks without requiring
+//        explicit coordination.
+// [^11]: See UpToRaftLogPosition in AdmittedRaftLogEntries.
 //
 // TODO(irfansharif): These descriptions are too high-level, imprecise and
 // possibly wrong. Fix that. After implementing these interfaces and integrating

pkg/server/node.go

@@ -907,7 +907,7 @@ func (n *Node) GetPebbleMetrics() []admission.StoreMetrics {
 			diskStats = s
 		}
 		metrics = append(metrics, admission.StoreMetrics{
-			StoreID:         int32(store.StoreID()),
+			StoreID:         store.StoreID(),
 			Metrics:         m.Metrics,
 			WriteStallCount: m.WriteStallCount,
 			DiskStats:       diskStats})

pkg/util/admission/BUILD.bazel

@@ -17,6 +17,7 @@ go_library(
         "scheduler_latency_listener.go",
         "sql_cpu_overload_indicator.go",
         "store_token_estimation.go",
+        "testing_knobs.go",
         "tokens_linear_model.go",
         "work_queue.go",
     ],
@@ -53,6 +54,7 @@ go_test(
         "elastic_cpu_work_queue_test.go",
         "granter_test.go",
         "io_load_listener_test.go",
+        "replicated_write_admission_test.go",
         "scheduler_latency_listener_test.go",
         "store_token_estimation_test.go",
         "tokens_linear_model_test.go",
@@ -67,6 +69,7 @@ go_test(
         "//pkg/testutils/datapathutils",
         "//pkg/testutils/echotest",
         "//pkg/util/admission/admissionpb",
+        "//pkg/util/humanizeutil",
         "//pkg/util/leaktest",
         "//pkg/util/log",
         "//pkg/util/metric",

pkg/util/admission/admission.go

@@ -522,8 +522,9 @@ func workKindString(workKind WorkKind) string {
 // all of these when StoreWorkQueue.AdmittedWorkDone is called, so that these
 // cumulative values are mutually consistent.
 type storeAdmissionStats struct {
-	// Total requests that called AdmittedWorkDone or BypassedWorkDone.
-	admittedCount uint64
+	// Total requests that called {Admitted,Bypassed}WorkDone, or in the case of
+	// replicated writes, the requests that called Admit.
+	workCount uint64
 	// Sum of StoreWorkDoneInfo.WriteBytes.
 	//
 	// TODO(sumeer): writeAccountedBytes and ingestedAccountedBytes are not
@@ -548,7 +549,7 @@ type storeAdmissionStats struct {
 	// (e.g. for logging).
 	aux struct {
 		// These bypassed numbers are already included in the corresponding
-		// {admittedCount, writeAccountedBytes, ingestedAccountedBytes}.
+		// {workCount, writeAccountedBytes, ingestedAccountedBytes}.
 		bypassedCount                  uint64
 		writeBypassedAccountedBytes    uint64
 		ingestedBypassedAccountedBytes uint64

pkg/util/admission/elastic_cpu_work_queue.go

@@ -81,7 +81,7 @@ func (e *ElasticCPUWorkQueue) Admit(
 	if duration > MaxElasticCPUDuration {
 		duration = MaxElasticCPUDuration
 	}
-	info.requestedCount = duration.Nanoseconds()
+	info.RequestedCount = duration.Nanoseconds()
 	enabled, err := e.workQueue.Admit(ctx, info)
 	if err != nil {
 		return nil, err

pkg/util/admission/elastic_cpu_work_queue_test.go

@@ -161,7 +161,7 @@ func (t *testElasticCPUInternalWorkQueue) Admit(
 	_ context.Context, info WorkInfo,
 ) (enabled bool, err error) {
 	if !t.disabled {
-		t.buf.WriteString(fmt.Sprintf("admitted=%s ", time.Duration(info.requestedCount)))
+		t.buf.WriteString(fmt.Sprintf("admitted=%s ", time.Duration(info.RequestedCount)))
 	}
 	return !t.disabled, nil
 }

pkg/util/admission/grant_coordinator.go

@@ -113,7 +113,7 @@ func (sgc *StoreGrantCoordinators) SetPebbleMetricsProvider(
 						if unsafeGc, ok := sgc.gcMap.Load(int64(m.StoreID)); ok {
 							gc := (*GrantCoordinator)(unsafeGc)
 							gc.pebbleMetricsTick(ctx, m)
-							iotc.UpdateIOThreshold(roachpb.StoreID(m.StoreID), gc.ioLoadListener.ioThreshold)
+							iotc.UpdateIOThreshold(m.StoreID, gc.ioLoadListener.ioThreshold)
 						} else {
 							log.Warningf(ctx,
 								"seeing metrics for unknown storeID %d", m.StoreID)
@@ -135,7 +135,7 @@ func (sgc *StoreGrantCoordinators) SetPebbleMetricsProvider(
 	}()
 }
 
-func (sgc *StoreGrantCoordinators) initGrantCoordinator(storeID int32) *GrantCoordinator {
+func (sgc *StoreGrantCoordinators) initGrantCoordinator(storeID roachpb.StoreID) *GrantCoordinator {
 	coord := &GrantCoordinator{
 		settings:       sgc.settings,
 		useGrantChains: false,
@@ -168,7 +168,7 @@ func (sgc *StoreGrantCoordinators) initGrantCoordinator(storeID int32) *GrantCoo
 		},
 	}
 
-	storeReq := sgc.makeStoreRequesterFunc(sgc.ambientCtx, granters, sgc.settings, sgc.workQueueMetrics, opts)
+	storeReq := sgc.makeStoreRequesterFunc(sgc.ambientCtx, storeID, granters, sgc.settings, sgc.workQueueMetrics, opts, nil)
 	coord.queues[KVWork] = storeReq
 	requesters := storeReq.getRequesters()
 	kvg.regularRequester = requesters[admissionpb.RegularWorkClass]
@@ -336,8 +336,9 @@ type makeRequesterFunc func(
 	metrics *WorkQueueMetrics, opts workQueueOptions) requester
 
 type makeStoreRequesterFunc func(
-	_ log.AmbientContext, granters [admissionpb.NumWorkClasses]granterWithStoreWriteDone,
-	settings *cluster.Settings, metrics *WorkQueueMetrics, opts workQueueOptions) storeRequester
+	_ log.AmbientContext, storeID roachpb.StoreID, granters [admissionpb.NumWorkClasses]granterWithStoreWriteDone,
+	settings *cluster.Settings, metrics *WorkQueueMetrics, opts workQueueOptions, knobs *TestingKnobs,
+) storeRequester
 
 // NewGrantCoordinators constructs GrantCoordinators and WorkQueues for a
 // regular cluster node. Caller is responsible for:

pkg/util/admission/granter.go

@@ -594,7 +594,7 @@ type IOThresholdConsumer interface {
 
 // StoreMetrics are the metrics for a store.
 type StoreMetrics struct {
-	StoreID int32
+	StoreID roachpb.StoreID
 	*pebble.Metrics
 	WriteStallCount int64
 	// Optional.

pkg/util/admission/granter_test.go

@@ -109,8 +109,9 @@ func TestGranterBasic(t *testing.T) {
 			storeCoordinators := &StoreGrantCoordinators{
 				settings: settings,
 				makeStoreRequesterFunc: func(
-					ambientCtx log.AmbientContext, granters [admissionpb.NumWorkClasses]granterWithStoreWriteDone,
-					settings *cluster.Settings, metrics *WorkQueueMetrics, opts workQueueOptions) storeRequester {
+					ambientCtx log.AmbientContext, _ roachpb.StoreID, granters [admissionpb.NumWorkClasses]granterWithStoreWriteDone,
+					settings *cluster.Settings, metrics *WorkQueueMetrics, opts workQueueOptions, knobs *TestingKnobs,
+				) storeRequester {
 					makeTestRequester := func(wc admissionpb.WorkClass) *testRequester {
 						req := &testRequester{
 							workKind:               KVWork,
@@ -273,8 +274,8 @@ func TestStoreCoordinators(t *testing.T) {
 	opts := Options{
 		makeRequesterFunc: makeRequesterFunc,
 		makeStoreRequesterFunc: func(
-			ctx log.AmbientContext, granters [admissionpb.NumWorkClasses]granterWithStoreWriteDone,
-			settings *cluster.Settings, metrics *WorkQueueMetrics, opts workQueueOptions) storeRequester {
+			ctx log.AmbientContext, _ roachpb.StoreID, granters [admissionpb.NumWorkClasses]granterWithStoreWriteDone,
+			settings *cluster.Settings, metrics *WorkQueueMetrics, opts workQueueOptions, _ *TestingKnobs) storeRequester {
 			reqReg := makeRequesterFunc(ctx, KVWork, granters[admissionpb.RegularWorkClass], settings, metrics, opts)
 			reqElastic := makeRequesterFunc(ctx, KVWork, granters[admissionpb.ElasticWorkClass], settings, metrics, opts)
 			str := &storeTestRequester{}
@@ -444,7 +445,7 @@ func (m *testMetricsProvider) setMetricsForStores(stores []int32, metrics pebble
 	m.metrics = m.metrics[:0]
 	for _, s := range stores {
 		m.metrics = append(m.metrics, StoreMetrics{
-			StoreID: s,
+			StoreID: roachpb.StoreID(s),
 			Metrics: &metrics,
 		})
 	}

pkg/util/admission/io_load_listener.go

@@ -15,6 +15,7 @@ import (
 	"math"
 	"time"
 
+	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/settings"
 	"github.com/cockroachdb/cockroach/pkg/settings/cluster"
 	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
@@ -151,7 +152,7 @@ const l0SubLevelCountOverloadThreshold = 20
 //     storeWriteDone), the tokens are adjusted differently for the
 //     flush/compaction L0 tokens and for the "disk bandwidth" tokens.
 type ioLoadListener struct {
-	storeID     int32
+	storeID     roachpb.StoreID
 	settings    *cluster.Settings
 	kvRequester storeRequester
 	kvGranter   granterWithIOTokens

pkg/util/admission/io_load_listener_test.go

@@ -67,11 +67,11 @@ func TestIOLoadListener(t *testing.T) {
 
 			case "prep-admission-stats":
 				req.stats = storeAdmissionStats{
-					admittedCount:          0,
+					workCount:              0,
 					writeAccountedBytes:    0,
 					ingestedAccountedBytes: 0,
 				}
-				d.ScanArgs(t, "admitted", &req.stats.admittedCount)
+				d.ScanArgs(t, "admitted", &req.stats.workCount)
 				if d.HasArg("write-bytes") {
 					d.ScanArgs(t, "write-bytes", &req.stats.writeAccountedBytes)
 				}
@@ -290,7 +290,7 @@ func TestBadIOLoadListenerStats(t *testing.T) {
 		d.BytesRead = rand.Uint64()
 		d.BytesWritten = rand.Uint64()
 		d.ProvisionedBandwidth = 1 << 20
-		req.stats.admittedCount = rand.Uint64()
+		req.stats.workCount = rand.Uint64()
 		req.stats.writeAccountedBytes = rand.Uint64()
 		req.stats.ingestedAccountedBytes = rand.Uint64()
 		req.stats.statsToIgnore.Bytes = rand.Uint64()