Merge #64734 #65029

64734: bulk: stop splitting on range flushes during RESTORE r=dt a=pbardea

The SST batcher maintains logic to preemptively split a range when it
sees the first size flush. However, we only want to do this when the SST
batcher is enabled to split and scatter ranges after consuming a certain
amount of data (which is indicated by setting the SplitAfter setting).

SplitAfter is enabled by default when the SST batcher is used by the
bulk adder (e.g. in IMPORT and backfills). In other cases (ie RESTORE),
spans are already pre-split so this splitting and scattering is
detrimental to RESTORE performance (range leaseholders get moved after
they went through the split and scattering process).

Release note: None

65029: bazel: handle `mockgen` in Bazel sandbox r=rail a=rickystewart

There's an open-source library to do this for us, and it works well
enough, so I just pull it in and use it everywhere we call into
`mockgen`.

Now there's only one un-bazelfied `go:generate` in-tree:
`pkg/server/api_v2.go` :)

Release note: None

Co-authored-by: Paul Bardea <[email protected]>
Co-authored-by: Ricky Stewart <[email protected]>

BUILD.bazel

@@ -75,6 +75,7 @@ load("@bazel_gazelle//:def.bzl", "gazelle")
 # gazelle:exclude **/*.pb.gw.go
 # gazelle:exclude **/*_interval_btree.go
 # gazelle:exclude **/*_interval_btree_test.go
+# gazelle:exclude **/mocks_generated.go
 # gazelle:exclude pkg/sql/parser/sql.go
 # gazelle:exclude pkg/sql/parser/helpmap_test.go
 # gazelle:exclude pkg/sql/parser/help_messages.go

WORKSPACE

@@ -178,3 +178,12 @@ rules_foreign_cc_dependencies()
 load("//build:toolchains/REPOSITORIES.bzl", "toolchain_dependencies")
 
 toolchain_dependencies()
+
+http_archive(
+    name = "bazel_gomock",
+    sha256 = "692421b0c5e04ae4bc0bfff42fb1ce8671fe68daee2b8d8ea94657bb1fcddc0a",
+    strip_prefix = "bazel_gomock-fde78c91cf1783cc1e33ba278922ba67a6ee2a84",
+    urls = [
+        "https://github.com/jmhodges/bazel_gomock/archive/fde78c91cf1783cc1e33ba278922ba67a6ee2a84.tar.gz",
+    ],
+)

pkg/ccl/sqlproxyccl/admitter/BUILD.bazel

@@ -1,19 +1,29 @@
+load("@bazel_gomock//:gomock.bzl", "gomock")
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
+# gazelle:exclude service.go
+
+go_library(
+    name = "service",
+    srcs = ["service.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/admitter",
+    visibility = ["//visibility:private"],
+)
+
 go_library(
     name = "admitter",
     srcs = [
         "local.go",
-        "mocks_generated.go",
-        "service.go",
+        ":mocks_admitter",  # keep
     ],
+    embed = [":service"],  # keep
     importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/admitter",
     visibility = ["//visibility:public"],
     deps = [
         "//pkg/util/log",
         "//pkg/util/syncutil",
         "@com_github_cockroachdb_errors//:errors",
-        "@com_github_golang_mock//gomock",
+        "@com_github_golang_mock//gomock",  # keep
     ],
 )
 
@@ -23,3 +33,11 @@ go_test(
     embed = [":admitter"],
     deps = ["@com_github_stretchr_testify//require"],
 )
+
+gomock(
+    name = "mocks_admitter",
+    out = "mocks_generated.go",
+    interfaces = ["Service"],
+    library = ":service",
+    package = "admitter",
+)

pkg/ccl/sqlproxyccl/denylist/BUILD.bazel

@@ -1,12 +1,22 @@
+load("@bazel_gomock//:gomock.bzl", "gomock")
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
+# gazelle:exclude service.go
+
+go_library(
+    name = "service",
+    srcs = ["service.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/denylist",
+    visibility = ["//visibility:private"],
+)
+
 go_library(
     name = "denylist",
     srcs = [
         "local_file.go",
-        "mocks_generated.go",
-        "service.go",
+        ":mocks_denylist",  # keep
     ],
+    embed = [":service"],  # keep
     importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/denylist",
     visibility = ["//visibility:public"],
     deps = [
@@ -15,7 +25,7 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
-        "@com_github_golang_mock//gomock",
+        "@com_github_golang_mock//gomock",  # keep
         "@com_github_spf13_viper//:viper",
     ],
 )
@@ -29,3 +39,12 @@ go_test(
         "@com_github_stretchr_testify//require",
     ],
 )
+
+gomock(
+    name = "mocks_denylist",
+    out = "mocks_generated.go",
+    interfaces = ["Service"],
+    library = ":service",
+    package = "denylist",
+    self_package = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/denylist",
+)

pkg/ccl/sqlproxyccl/tenant/BUILD.bazel

@@ -1,6 +1,7 @@
-load("@rules_proto//proto:defs.bzl", "proto_library")
+load("@bazel_gomock//:gomock.bzl", "gomock")
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
+load("@rules_proto//proto:defs.bzl", "proto_library")
 
 proto_library(
     name = "tenant_proto",
@@ -24,8 +25,8 @@ go_library(
     srcs = [
         "directory.go",
         "entry.go",
-        "mocks_generated.go",
         "test_directory_svr.go",
+        ":mocks_tenant",  # keep
     ],
     embed = [":tenant_go_proto"],
     importpath = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant",
@@ -39,9 +40,9 @@ go_library(
         "//pkg/util/timeutil",
         "@com_github_cockroachdb_errors//:errors",
         "@com_github_cockroachdb_logtags//:logtags",
-        "@com_github_golang_mock//gomock",
+        "@com_github_golang_mock//gomock",  # keep
         "@org_golang_google_grpc//:go_default_library",
-        "@org_golang_google_grpc//metadata",
+        "@org_golang_google_grpc//metadata",  # keep
         "@org_golang_google_grpc//status",
     ],
 )
@@ -73,3 +74,15 @@ go_test(
         "@org_golang_google_grpc//:go_default_library",
     ],
 )
+
+gomock(
+    name = "mocks_tenant",
+    out = "mocks_generated.go",
+    interfaces = [
+        "DirectoryClient",
+        "Directory_WatchEndpointsClient",
+    ],
+    library = ":tenant_go_proto",
+    package = "tenant",
+    self_package = "github.com/cockroachdb/cockroach/pkg/ccl/sqlproxyccl/tenant",
+)

pkg/kv/bulk/buffering_adder.go

@@ -94,7 +94,7 @@ func MakeBulkAdder(
 		// splitting _before_ hitting max reduces chance of auto-splitting after the
 		// range is full and is more expensive to split/move.
 		opts.SplitAndScatterAfter = func() int64 { return 48 << 20 }
-	} else if opts.SplitAndScatterAfter() == -1 {
+	} else if opts.SplitAndScatterAfter() == kvserverbase.DisableExplicitSplits {
 		opts.SplitAndScatterAfter = nil
 	}
 

pkg/kv/bulk/sst_batcher.go

@@ -281,7 +281,11 @@ func (b *SSTBatcher) doFlush(ctx context.Context, reason int, nextKey roachpb.Ke
 		// non-overlapping keyspace this split partitions off "our" target space for
 		// future splitting/scattering, while if they don't, doing this only once
 		// minimizes impact on other adders (e.g. causing extra SST splitting).
-		if b.flushCounts.total == 1 {
+		//
+		// We only do this splitting if the caller expects the sst_batcher to
+		// split and scatter the data as it ingests it (which is the case when
+		// splitAfter) is set.
+		if b.flushCounts.total == 1 && b.splitAfter != nil {
 			if splitAt, err := keys.EnsureSafeSplitKey(start); err != nil {
 				log.Warningf(ctx, "%v", err)
 			} else {

pkg/security/certmgr/BUILD.bazel

@@ -1,13 +1,23 @@
+load("@bazel_gomock//:gomock.bzl", "gomock")
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 
+# gazelle:exclude cert.go
+
+go_library(
+    name = "certlib",
+    srcs = ["cert.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/security/certmgr",
+    visibility = ["//visibility:private"],
+)
+
 go_library(
     name = "certmgr",
     srcs = [
-        "cert.go",
         "cert_manager.go",
         "file_cert.go",
-        "mocks_generated.go",
+        ":mocks_certmgr",  # keep
     ],
+    embed = [":certlib"],  # keep
     importpath = "github.com/cockroachdb/cockroach/pkg/security/certmgr",
     visibility = ["//visibility:public"],
     deps = [
@@ -16,7 +26,7 @@ go_library(
         "//pkg/util/syncutil",
         "//pkg/util/sysutil",
         "@com_github_cockroachdb_errors//:errors",
-        "@com_github_golang_mock//gomock",
+        "@com_github_golang_mock//gomock",  # keep
     ],
 )
 
@@ -39,3 +49,11 @@ go_test(
         "@org_golang_x_sys//unix",
     ],
 )
+
+gomock(
+    name = "mocks_certmgr",
+    out = "mocks_generated.go",
+    interfaces = ["Cert"],
+    library = ":certlib",
+    package = "certmgr",
+)