kvserver: return NodeUnavailable error during range ID allocation

During a split, the request to generate a new range ID for the RHS can
race with shutting down the node where the underlying increment is
evaluated. Previously, in this case, the range ID allocator would
return a custom error (`could not allocate ID; system is draining`),
which is not retried by the `DistSender`.

This commit replaces this error with a `NodeUnavailableError`, which is
how this type of race is handled in other places as well.

Informs: #64828

Release note: None

Author: miraradeva

pkg/kv/kvserver/idalloc/BUILD.bazel

@@ -8,6 +8,7 @@ go_library(
     deps = [
         "//pkg/base",
         "//pkg/kv",
+        "//pkg/kv/kvpb",
         "//pkg/roachpb",
         "//pkg/util/log",
         "//pkg/util/retry",
@@ -28,6 +29,7 @@ go_test(
     deps = [
         "//pkg/keys",
         "//pkg/kv/kvclient/kvcoord",
+        "//pkg/kv/kvpb",
         "//pkg/roachpb",
         "//pkg/security/securityassets",
         "//pkg/security/securitytest",

pkg/kv/kvserver/idalloc/id_alloc.go

@@ -11,6 +11,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/base"
 	"github.com/cockroachdb/cockroach/pkg/kv"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	"github.com/cockroachdb/cockroach/pkg/util/log"
 	"github.com/cockroachdb/cockroach/pkg/util/retry"
@@ -84,7 +85,7 @@ func (ia *Allocator) Allocate(ctx context.Context) (int64, error) {
 	case id := <-ia.ids:
 		// when the channel is closed, the zero value is returned.
 		if id == 0 {
-			return id, errors.Errorf("could not allocate ID; system is draining")
+			return id, &kvpb.NodeUnavailableError{}
 		}
 		return id, nil
 	case <-ctx.Done():

pkg/kv/kvserver/idalloc/id_alloc_test.go

@@ -14,6 +14,7 @@ import (
 
 	"github.com/cockroachdb/cockroach/pkg/keys"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvclient/kvcoord"
+	"github.com/cockroachdb/cockroach/pkg/kv/kvpb"
 	"github.com/cockroachdb/cockroach/pkg/kv/kvserver/idalloc"
 	"github.com/cockroachdb/cockroach/pkg/roachpb"
 	// Import to set ZoneConfigHook.
@@ -232,9 +233,9 @@ func TestAllocateWithStopper(t *testing.T) {
 	s, idAlloc := newTestAllocator(t)
 	s.Stop() // not deferred.
 
-	if _, err := idAlloc.Allocate(context.Background()); !testutils.IsError(err, "system is draining") {
-		t.Errorf("unexpected error: %+v", err)
-	}
+	_, err := idAlloc.Allocate(context.Background())
+	require.Error(t, err)
+	require.ErrorIs(t, err, &kvpb.NodeUnavailableError{})
 }
 
 // TestLostWriteAssertion makes sure that the Allocator performs a best-effort