roachprod: add typegen code generator

Previously, when using the codec utility, all dynamic types needed to implement
the `codec.DynamicType` interface and register itself against the codec package
manually.

This change provides a code generator `typegen` for generating type registration
source files. Users of the codec utility can now simply add `typegen:reg` in a
comment above the concrete implementations of dynamic types.

Then it only requires adding the `genrule` in the bazel file of the package that
wants to generate the type registrations, and the type registration source files
will be generated automatically.

Informs: #149451, #151461

Epic: None
Release note: None

Author: herkolategan

pkg/BUILD.bazel

@@ -1699,6 +1699,8 @@ GO_TARGETS = [
     "//pkg/roachprod/prometheus:prometheus_test",
     "//pkg/roachprod/promhelperclient:promhelperclient",
     "//pkg/roachprod/promhelperclient:promhelperclient_test",
+    "//pkg/roachprod/roachprodutil/codec/cmd/typegen:typegen",
+    "//pkg/roachprod/roachprodutil/codec/cmd/typegen:typegen_lib",
     "//pkg/roachprod/roachprodutil/codec:codec",
     "//pkg/roachprod/roachprodutil/codec:codec_test",
     "//pkg/roachprod/roachprodutil:roachprodutil",

pkg/roachprod/roachprodutil/codec/cmd/typegen/BUILD.bazel

@@ -0,0 +1,18 @@
+load("@io_bazel_rules_go//go:def.bzl", "go_binary", "go_library")
+
+go_library(
+    name = "typegen_lib",
+    srcs = ["main.go"],
+    importpath = "github.com/cockroachdb/cockroach/pkg/roachprod/roachprodutil/codec/cmd/typegen",
+    visibility = ["//visibility:public"],
+    deps = [
+        "//pkg/cli/exit",
+        "@com_github_cockroachdb_errors//:errors",
+    ],
+)
+
+go_binary(
+    name = "typegen",
+    embed = [":typegen_lib"],
+    visibility = ["//visibility:public"],
+)

pkg/roachprod/roachprodutil/codec/cmd/typegen/main.go

@@ -0,0 +1,153 @@
+// Copyright 2025 The Cockroach Authors.
+//
+// Use of this software is governed by the CockroachDB Software License
+// included in the /LICENSE file.
+
+package main
+
+import (
+	"flag"
+	"fmt"
+	"go/ast"
+	"go/parser"
+	"go/token"
+	"os"
+	"path/filepath"
+	"sort"
+	"strings"
+
+	"github.com/cockroachdb/cockroach/pkg/cli/exit"
+	"github.com/cockroachdb/errors"
+)
+
+var (
+	outputFile = flag.String("output", "", "Output file; defaults to stdout.")
+)
+
+func usage() {
+	fmt.Println(`usage: typegen [flags] <source> [<type> ...]
+
+This program generates an implementation of the [codec.DynamicType] interface
+for each specified type present in the source file. It will also create an
+'init function' that registers the types with the codec package.
+
+Any type documented with the "// typegen:reg" comment above its definition will
+also be registered.
+
+Flags:`)
+	flag.PrintDefaults()
+}
+
+func exitOnError(err error) {
+	if err == nil {
+		return
+	}
+	_, _ = fmt.Fprintf(os.Stderr, "%s\n", err)
+	exit.WithCode(exit.UnspecifiedError())
+}
+
+func main() {
+	// Flag handling section
+	flag.Usage = usage
+	flag.Parse()
+
+	if flag.NArg() < 1 {
+		usage()
+		exit.WithCode(exit.UnspecifiedError())
+	}
+	source := flag.Arg(0)
+	verifyTypes := make(map[string]struct{})
+
+	for _, t := range flag.Args()[1:] {
+		verifyTypes[t] = struct{}{}
+	}
+
+	// Parse the source file.
+	fs := token.NewFileSet()
+	file, err := parser.ParseFile(fs, source, nil, parser.ParseComments)
+	exitOnError(errors.Wrapf(err, "parsing %s", source))
+
+	// Find all types that need to be registered.
+	types := make(map[string]struct{})
+	for _, decl := range file.Decls {
+		gen, ok := decl.(*ast.GenDecl)
+		if !ok || gen.Tok != token.TYPE {
+			continue
+		}
+		for _, spec := range gen.Specs {
+			ts, ok := spec.(*ast.TypeSpec)
+			if !ok {
+				continue
+			}
+
+			typeName := ts.Name.Name
+			if gen.Doc != nil {
+				text := gen.Doc.Text()
+				if strings.Contains(text, "typegen:reg") {
+					types[typeName] = struct{}{}
+				}
+			}
+
+			if _, ok := verifyTypes[typeName]; ok {
+				delete(verifyTypes, typeName)
+				types[typeName] = struct{}{}
+				continue
+			}
+		}
+	}
+
+	// Validate that types passed in via the command line were found.
+	if len(verifyTypes) > 0 {
+		var missing []string
+		for t := range verifyTypes {
+			missing = append(missing, t)
+		}
+		exitOnError(errors.Newf("the following types were not found in %s: %v", source, missing))
+	}
+
+	// Fail if no types were found.
+	if len(types) == 0 {
+		exitOnError(errors.Newf("no types found in %s", source))
+	}
+
+	var typeKeys []string
+	for typeName := range types {
+		typeKeys = append(typeKeys, typeName)
+	}
+	sort.Strings(typeKeys)
+
+	// Generate the type registration code.
+	w := os.Stdout
+	if *outputFile != "" {
+		w, err = os.Create(*outputFile)
+		exitOnError(errors.Wrapf(err, "creating %s", *outputFile))
+		defer w.Close()
+	}
+	p := func(format string, args ...interface{}) {
+		_, _ = fmt.Fprintf(w, format, args...)
+	}
+
+	pkg := file.Name.Name
+
+	p("// Code generated by typegen; DO NOT EDIT.\n")
+	p("// Generated from: %s\n\n", filepath.Base(source))
+	p("package %s\n\n", pkg)
+
+	prefixCall := ""
+	if pkg != "codec" {
+		prefixCall = "codec."
+		p("import \"github.com/cockroachdb/cockroach/pkg/roachprod/roachprodutil/codec\"\n\n")
+	}
+
+	p("func init() {\n")
+	for _, t := range typeKeys {
+		p("	%sRegister(new(%s))\n", prefixCall, t)
+	}
+	p("}\n")
+
+	for _, t := range typeKeys {
+		p("\nfunc (t %s) GetTypeName() %sTypeName {\n", t, prefixCall)
+		p("	return %sResolveTypeName(t)\n", prefixCall)
+		p("}\n")
+	}
+}