roachpb: stop using reflection to generate file

This reflection ended up causing serious build pain and represented a non-zero
portion of the critical path time of our build. AST inspection works just as
well as reflection and means we now only depend on the generated proto files.

Release note: None

.github/CODEOWNERS

@@ -229,6 +229,7 @@
 /pkg/roachpb/dep_test.go     @cockroachdb/dev-inf
 /pkg/roachpb/error*          @cockroachdb/kv-prs
 /pkg/roachpb/gen             @cockroachdb/dev-inf
+/pkg/roachpb/gen.bzl         @cockroachdb/dev-inf
 /pkg/roachpb/app*            @cockroachdb/sql-observability
 /pkg/roachpb/index*          @cockroachdb/sql-observability
 /pkg/roachpb/internal*       @cockroachdb/kv-prs

BUILD.bazel

@@ -57,16 +57,6 @@ exports_files([
 # gazelle:resolve go google.golang.org/genproto/googleapis/pubsub/v1 @org_golang_google_genproto//googleapis/pubsub/v1:pubsub
 # gazelle:resolve go google.golang.org/genproto/googleapis/cloud/kms/v1 @org_golang_google_genproto//googleapis/cloud/kms/v1:kms
 
-# These packages use github.com/golang/mock to generate mocks. The target
-# generating mocks for a given package necessarily depends on the package's
-# go_library target, which then prevents us from including those mock symbols
-# as part of the package's go_library target. We consequently export a second
-# target, embedding the first but also including the mock symbols. We need to
-# tell gazelle to resolve to this second target instead. See
-# pkg/kv/kvclient/rangefeed/BUILD.bazel for an annotated example.
-#
-# gazelle:resolve go github.com/cockroachdb/cockroach/pkg/roachpb //pkg/roachpb
-
 # See pkg/roachpb/gen/BUILD.bazel for more details.
 #
 # gazelle:resolve proto go roachpb/api.proto //pkg/roachpb

build/bazelutil/check.sh

@@ -18,7 +18,7 @@ pkg/kv/kvclient/rangefeed/rangefeed.go://go:generate mockgen -destination=mocks_
 pkg/kv/kvserver/concurrency/lock_table.go://go:generate ../../../util/interval/generic/gen.sh *lockState concurrency
 pkg/kv/kvserver/spanlatch/manager.go://go:generate ../../../util/interval/generic/gen.sh *latch spanlatch
 pkg/roachpb/api.go://go:generate mockgen -package=roachpbmock -destination=roachpbmock/mocks_generated.go . InternalClient,Internal_RangeFeedClient
-pkg/roachpb/batch.go://go:generate go run -tags gen-batch gen/main.go
+pkg/roachpb/batch.go://go:generate go run gen/main.go --filename batch_generated.go *.pb.go
 pkg/security/certmgr/cert.go://go:generate mockgen -package=certmgr -destination=mocks_generated_test.go . Cert
 pkg/security/securitytest/securitytest.go://go:generate go-bindata -mode 0600 -modtime 1400000000 -pkg securitytest -o embedded.go -ignore README.md -ignore regenerate.sh test_certs
 pkg/security/securitytest/securitytest.go://go:generate gofmt -s -w embedded.go

pkg/cmd/bazci/bazci.go

@@ -201,7 +201,7 @@ func getBuildInfo(args parsedArgs) (buildInfo, error) {
 		switch targetKind {
 		case "cmake":
 			ret.cmakeTargets = append(ret.cmakeTargets, fullTarget)
-		case "genrule":
+		case "genrule", "batch_gen":
 			ret.genruleTargets = append(ret.genruleTargets, fullTarget)
 		case "go_binary":
 			ret.goBinaries = append(ret.goBinaries, fullTarget)

pkg/roachpb/BUILD.bazel

@@ -4,13 +4,8 @@ load("@rules_proto//proto:defs.bzl", "proto_library")
 load("@io_bazel_rules_go//proto:def.bzl", "go_proto_library")
 load("@io_bazel_rules_go//go:def.bzl", "go_library", "go_test")
 load("//build:STRINGER.bzl", "stringer")
+load(":gen.bzl", "batch_gen")
 
-# NB: If you've come here to resolve a build failure due to a missing
-# dependency, the dependency may be missing from the "bootstrap" target rather
-# than the "roachpb" target. Some input source files are duplicated between the
-# two targets, and Gazelle only infers dependencies for "roachpb" rather than
-# "bootstrap". See pkg/roachpb/gen/BUILD.bazel for more detail on why the
-# "bootstrap" target exists.
 go_library(
     name = "roachpb",
     srcs = [
@@ -64,59 +59,6 @@ go_library(
     ],
 )
 
-# See the BUILD.bazel file in pkg/roachpb/gen for more details. We need to
-# partition the dependencies here to avoid cyclical structures.
-#
-# keep
-go_library(
-    name = "bootstrap",
-    srcs = [
-        "api.go",
-        "app_stats.go",
-        "batch.go",
-        "batch_generated.go",
-        "data.go",
-        "errors.go",
-        "index_usage_stats.go",
-        "metadata.go",
-        "metadata_replicas.go",
-        "method.go",
-        "span_config.go",
-        "tenant.go",
-        "version.go",
-    ],
-    embed = [":roachpb_go_proto"],
-    importpath = "github.com/cockroachdb/cockroach/pkg/roachpb",
-    visibility = [":__subpackages__"],
-    deps = [
-        "//pkg/geo",
-        "//pkg/geo/geopb",
-        "//pkg/kv/kvserver/concurrency/lock",
-        "//pkg/storage/enginepb",
-        "//pkg/util",
-        "//pkg/util/bitarray",
-        "//pkg/util/caller",
-        "//pkg/util/duration",
-        "//pkg/util/encoding",
-        "//pkg/util/hlc",
-        "//pkg/util/humanizeutil",
-        "//pkg/util/interval",
-        "//pkg/util/log",
-        "//pkg/util/protoutil",
-        "//pkg/util/timetz",
-        "//pkg/util/tracing",
-        "//pkg/util/uuid",
-        "@com_github_aws_aws_sdk_go//aws",
-        "@com_github_aws_aws_sdk_go//aws/credentials",
-        "@com_github_cockroachdb_apd_v3//:apd",
-        "@com_github_cockroachdb_errors//:errors",
-        "@com_github_cockroachdb_errors//errorspb",
-        "@com_github_cockroachdb_errors//extgrpc",
-        "@com_github_cockroachdb_redact//:redact",
-        "@io_etcd_go_etcd_raft_v3//raftpb",
-    ],
-)
-
 go_test(
     name = "roachpb_test",
     size = "small",
@@ -259,11 +201,8 @@ stringer(
     typ = "ErrorDetailType",
 )
 
-genrule(
+batch_gen(
     name = "gen-batch-generated",
-    outs = ["batch_generated-gen.go"],
-    cmd = """
-     $(location //pkg/roachpb/gen) --filename=$(location batch_generated-gen.go)
-       """,
-    exec_tools = ["//pkg/roachpb/gen"],
+    src = ":roachpb_go_proto",
+    out = "batch_generated-gen.go",
 )

pkg/roachpb/batch.go

@@ -23,7 +23,7 @@ import (
 	"github.com/cockroachdb/redact"
 )
 
-//go:generate go run -tags gen-batch gen/main.go
+//go:generate go run gen/main.go --filename batch_generated.go *.pb.go
 
 // WriteTimestamp returns the timestamps at which this request is writing. For
 // non-transactional requests, this is the same as the read timestamp. For

pkg/roachpb/gen.bzl

@@ -0,0 +1,21 @@
+load("@io_bazel_rules_go//go:def.bzl", "GoSource")
+
+def _batch_gen_impl(ctx):
+    srcs = [src for src in ctx.attr.src[GoSource].srcs]
+    print(ctx.outputs.out)
+    ctx.actions.run(
+        outputs = [ctx.outputs.out],
+        inputs = srcs,
+        executable = ctx.executable._tool,
+        arguments = ["--filename", ctx.outputs.out.path] + [src.path for src in srcs],
+    )
+    return [DefaultInfo(files = depset([ctx.outputs.out])),]
+
+batch_gen = rule(
+   implementation = _batch_gen_impl,
+   attrs = {
+       "out": attr.output(mandatory = True),
+       "src": attr.label(providers = [GoSource]),
+       "_tool": attr.label(default = "//pkg/roachpb/gen", executable = True, cfg = "exec"),
+   },
+)

pkg/roachpb/gen/BUILD.bazel

@@ -1,31 +1,11 @@
 load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
 
-# gen, as written, depends on symbols defined within //pkg/roachpb. gen is used
-# to generated a file[1] that in turn is depended on to compile //pkg/roachpb.
-# To break this circular structure, we define a separate target[2] within
-# pkg/roachpb that enlists the pre-generated file in the source tree[3] in
-# order to generate newer revisions within the sandbox. For another example of
-# this pattern, see what we do in langgen[4].
-#
-# All dependents of //pkg/roachpb will need to depend on the file generated
-# within the Bazel sandbox, which is achieved by [5]. Conversely, to instruct
-# gazelle/Bazel to resolve gen's import of //pkg/roachpb appropriately, we add
-# the resolve directive below.
-#
-# [1]: //pkg/roachpb/gen-batch-generated
-# [2]: //pkg/roachpb:bootstrap
-# [3]: //pkg/roachpb:batch_generated.go
-# [4]: //pkg/sql/opt/optgen/lang:bootstrap
-# [5]: See the "gazelle:resolve" directive in the top-level BUILD.bazel.
-#
-# gazelle:resolve go github.com/cockroachdb/cockroach/pkg/roachpb //pkg/roachpb:bootstrap
-
 go_library(
     name = "gen_lib",
     srcs = ["main.go"],
     importpath = "github.com/cockroachdb/cockroach/pkg/roachpb/gen",
     visibility = ["//visibility:private"],
-    deps = ["//pkg/roachpb:bootstrap"],
+    deps = ["@org_golang_x_tools//go/ast/inspector"],
 )
 
 go_binary(

pkg/roachpb/gen/main.go

@@ -8,20 +8,20 @@
 // by the Apache License, Version 2.0, included in the file
 // licenses/APL.txt.
 
-// This file generates batch_generated.go. It can be run via:
-//    go run -tags gen-batch gen/main.go
-
 package main
 
 import (
 	"flag"
 	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
 	"io"
 	"os"
-	"reflect"
+	"path/filepath"
 	"strings"
 
-	"github.com/cockroachdb/cockroach/pkg/roachpb"
+	"golang.org/x/tools/go/ast/inspector"
 )
 
 type variantInfo struct {
@@ -41,35 +41,34 @@ var reqVariants []variantInfo
 var resVariants []variantInfo
 var reqResVariantMapping map[variantInfo]variantInfo
 
-func initVariant(varInstance interface{}) variantInfo {
-	t := reflect.TypeOf(varInstance)
-	f := t.Elem().Field(0) // variants always have 1 field
+func initVariant(ins *inspector.Inspector, varName string) variantInfo {
+	fieldName, msgName := findVariantField(ins, varName)
 	return variantInfo{
-		variantType: t.Elem().Name(),
-		variantName: f.Name,
-		msgType:     f.Type.Elem().Name(),
+		variantType: varName,
+		variantName: fieldName,
+		msgType:     msgName,
 	}
 }
 
-func initVariants() {
-	errVars := (&roachpb.ErrorDetail{}).XXX_OneofWrappers()
+func initVariants(ins *inspector.Inspector) {
+	errVars := findVariantTypes(ins, "ErrorDetail")
 	for _, v := range errVars {
-		errInfo := initVariant(v)
+		errInfo := initVariant(ins, v)
 		errVariants = append(errVariants, errInfo)
 	}
 
-	resVars := (&roachpb.ResponseUnion{}).XXX_OneofWrappers()
+	resVars := findVariantTypes(ins, "ResponseUnion")
 	resVarInfos := make(map[string]variantInfo, len(resVars))
 	for _, v := range resVars {
-		resInfo := initVariant(v)
+		resInfo := initVariant(ins, v)
 		resVariants = append(resVariants, resInfo)
 		resVarInfos[resInfo.variantName] = resInfo
 	}
 
-	reqVars := (&roachpb.RequestUnion{}).XXX_OneofWrappers()
+	reqVars := findVariantTypes(ins, "RequestUnion")
 	reqResVariantMapping = make(map[variantInfo]variantInfo, len(reqVars))
 	for _, v := range reqVars {
-		reqInfo := initVariant(v)
+		reqInfo := initVariant(ins, v)
 		reqVariants = append(reqVariants, reqInfo)
 
 		// The ResponseUnion variants match those in RequestUnion, with the
@@ -87,6 +86,95 @@ func initVariants() {
 	}
 }
 
+// findVariantTypes leverages the fact that the protobuf generations creates
+// a method called XXX_OneofWrappers for oneof message types. The body of that
+// method is always a single return expression which returns a []interface{}
+// where each of the element in the slice literal are (*<type>)(nil)
+// expressions which can be interpreted using reflection. We'll find that
+// method for the oneof type with the requested name and then fish out the
+// list of variant types from inside that ParenExpr in the CompositeLit
+// underneath that method.
+//
+// The code in question looks like the below snippet, where we would pull
+// "ErrorDetail_NotLeaseholder" one of the returned strings.
+//
+//  // XXX_OneofWrappers is for the internal use of the proto package.
+//  func (*ErrorDetail) XXX_OneofWrappers() []interface{} {
+//  return []interface{}{
+//    (*ErrorDetail_NotLeaseHolder)(nil),
+//    ...
+//
+func findVariantTypes(ins *inspector.Inspector, oneofName string) []string {
+	var variants []string
+	var inFunc bool
+	var inParen bool
+	ins.Nodes([]ast.Node{
+		(*ast.FuncDecl)(nil),
+		(*ast.ParenExpr)(nil),
+		(*ast.Ident)(nil),
+	}, func(node ast.Node, push bool) (proceed bool) {
+		switch n := node.(type) {
+		case *ast.FuncDecl:
+			if n.Name.Name != "XXX_OneofWrappers" {
+				return false
+			}
+			se, ok := n.Recv.List[0].Type.(*ast.StarExpr)
+			if !ok {
+				return false
+			}
+			if se.X.(*ast.Ident).Name != oneofName {
+				return false
+			}
+			inFunc = push
+			return true
+		case *ast.ParenExpr:
+			inParen = push && inFunc
+			return inParen
+		case *ast.Ident:
+			if inParen {
+				variants = append(variants, n.Name)
+			}
+			return false
+		default:
+			return false
+		}
+	})
+	return variants
+}
+
+// findVariantField is used to find the field name and type name of the
+// variant with the name vType. Oneof variant types always have a single
+// field.
+//
+// The code in question looks like the below snippet, where we would return
+// ("NotLeaseHolder", "NotLeaseHolderError").
+//
+//  type ErrorDetail_NotLeaseHolder struct {
+//    NotLeaseHolder *NotLeaseHolderError
+//  }
+//
+func findVariantField(ins *inspector.Inspector, vType string) (fieldName, msgName string) {
+	ins.Preorder([]ast.Node{
+		(*ast.TypeSpec)(nil),
+	}, func(node ast.Node) {
+		n := node.(*ast.TypeSpec)
+		if n.Name.Name != vType {
+			return
+		}
+		st, ok := n.Type.(*ast.StructType)
+		if !ok {
+			return
+		}
+		if len(st.Fields.List) != 1 {
+			panic(fmt.Sprintf("type %v has %d fields", vType, len(st.Fields.List)))
+		}
+		f := st.Fields.List[0]
+		fieldName = f.Names[0].Name
+		msgName = f.Type.(*ast.StarExpr).X.(*ast.Ident).Name
+	})
+	return fieldName, msgName
+}
+
 func genGetInner(w io.Writer, unionName, variantName string, variants []variantInfo) {
 	fmt.Fprintf(w, `
 // GetInner returns the %[2]s contained in the union.
@@ -133,8 +221,24 @@ func (ru *%[1]s) MustSetInner(r %[2]s) {
 func main() {
 	var filenameFlag = flag.String("filename", "batch_generated.go", "filename path")
 	flag.Parse()
+	fileSet := token.NewFileSet()
+	var files []*ast.File
+	for _, arg := range flag.Args() {
+		expanded, err := filepath.Glob(arg)
+		if err != nil {
+			panic(fmt.Sprintf("failed to expand %s: %v", arg, err))
+		}
+		for _, fp := range expanded {
+			f, err := parser.ParseFile(fileSet, fp, nil, 0)
+			if err != nil {
+				panic(fmt.Sprintf("failed to parse %s: %v", arg, err))
+			}
+			files = append(files, f)
+		}
+	}
 
-	initVariants()
+	ins := inspector.New(files)
+	initVariants(ins)
 
 	f, err := os.Create(*filenameFlag)
 	if err != nil {