rac2: introduce stream token watcher interface

Introduce the `StreamTokenWatcher` interface, which will be used to
watch and be notified when there are tokens available.

Resolves: cockroachdb#128011
Release note: None

pkg/kv/kvserver/kvflowcontrol/rac2/store_stream.go

@@ -10,7 +10,12 @@
 
 package rac2
 
-import "github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol"
+import (
+	"context"
+
+	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/kvflowcontrol"
+	"github.com/cockroachdb/cockroach/pkg/util/admission/admissionpb"
+)
 
 // StreamTokenCounterProvider is the interface for retrieving token counters
 // for a given stream.
@@ -22,3 +27,36 @@ type StreamTokenCounterProvider interface {
 	// stream.
 	SendTokenCounterForStream(kvflowcontrol.Stream) TokenCounter
 }
+
+// StreamWatcherHandleID is a unique identifier for a handle that is
+// watching for available tokens on a stream.
+type StreamWatcherHandleID int64
+
+// StreamTokenWatcher is the interface for watching and waiting on available
+// tokens. The watcher registers a notification, which will be called when send
+// tokens are available for a given stream and work class.
+type StreamTokenWatcher interface {
+	// NotifyWhenAvailable queues up for tokens for the given token counter and
+	// work class. When tokens are available, the provided TokenGrantNotification
+	// is called. It is the caller's responsibility to call CancelHandle when
+	// tokens are no longer needed, or when the caller is done.
+	NotifyWhenAvailable(
+		TokenCounter,
+		admissionpb.WorkClass,
+		TokenGrantNotification,
+	) StreamWatcherHandleID
+	// UpdateHandle updates the given handle to watch the given work class,
+	// removing it from watching the existing work class, if the work class is
+	// different.
+	UpdateHandle(StreamWatcherHandleID, admissionpb.WorkClass)
+	// CancelHandle cancels the given handle, stopping it from being notified
+	// when tokens are available.
+	CancelHandle(StreamWatcherHandleID)
+}
+
+// TokenGrantNotification is an interface for notifying when tokens are
+// available.
+type TokenGrantNotification interface {
+	// Notify is called when tokens are available to be granted.
+	Notify(context.Context)
+}