github.com/aakash4dev/cometbft@v0.38.2/scripts/metricsgen/metricsgen.go (about) 1 // metricsgen is a code generation tool for creating constructors for CometBFT 2 // metrics types. 3 package main 4 5 import ( 6 "bytes" 7 "flag" 8 "fmt" 9 "go/ast" 10 "go/format" 11 "go/parser" 12 "go/token" 13 "go/types" 14 "io" 15 "io/fs" 16 "log" 17 "os" 18 "path" 19 "path/filepath" 20 "reflect" 21 "regexp" 22 "strconv" 23 "strings" 24 "text/template" 25 ) 26 27 func init() { 28 flag.Usage = func() { 29 fmt.Fprintf(os.Stderr, `Usage: %[1]s -struct <struct> 30 31 Generate constructors for the metrics type specified by -struct contained in 32 the current directory. The tool creates a new file in the current directory 33 containing the generated code. 34 35 Options: 36 `, filepath.Base(os.Args[0])) 37 flag.PrintDefaults() 38 } 39 } 40 41 const metricsPackageName = "github.com/go-kit/kit/metrics" 42 43 const ( 44 metricNameTag = "metrics_name" 45 labelsTag = "metrics_labels" 46 bucketTypeTag = "metrics_buckettype" 47 bucketSizeTag = "metrics_bucketsizes" 48 ) 49 50 var ( 51 dir = flag.String("dir", ".", "Path to the directory containing the target package") 52 strct = flag.String("struct", "Metrics", "Struct to parse for metrics") 53 ) 54 55 var bucketType = map[string]string{ 56 "exprange": "stdprometheus.ExponentialBucketsRange", 57 "exp": "stdprometheus.ExponentialBuckets", 58 "lin": "stdprometheus.LinearBuckets", 59 } 60 61 var tmpl = template.Must(template.New("tmpl").Parse(`// Code generated by metricsgen. DO NOT EDIT. 62 63 package {{ .Package }} 64 65 import ( 66 "github.com/go-kit/kit/metrics/discard" 67 prometheus "github.com/go-kit/kit/metrics/prometheus" 68 stdprometheus "github.com/prometheus/client_golang/prometheus" 69 ) 70 71 func PrometheusMetrics(namespace string, labelsAndValues...string) *Metrics { 72 labels := []string{} 73 for i := 0; i < len(labelsAndValues); i += 2 { 74 labels = append(labels, labelsAndValues[i]) 75 } 76 return &Metrics{ 77 {{ range $metric := .ParsedMetrics }} 78 {{- $metric.FieldName }}: prometheus.New{{ $metric.TypeName }}From(stdprometheus.{{$metric.TypeName }}Opts{ 79 Namespace: namespace, 80 Subsystem: MetricsSubsystem, 81 Name: "{{$metric.MetricName }}", 82 Help: "{{ $metric.Description }}", 83 {{ if ne $metric.HistogramOptions.BucketType "" }} 84 Buckets: {{ $metric.HistogramOptions.BucketType }}({{ $metric.HistogramOptions.BucketSizes }}), 85 {{ else if ne $metric.HistogramOptions.BucketSizes "" }} 86 Buckets: []float64{ {{ $metric.HistogramOptions.BucketSizes }} }, 87 {{ end }} 88 {{- if eq (len $metric.Labels) 0 }} 89 }, labels).With(labelsAndValues...), 90 {{ else }} 91 }, append(labels, {{$metric.Labels}})).With(labelsAndValues...), 92 {{ end }} 93 {{- end }} 94 } 95 } 96 97 98 func NopMetrics() *Metrics { 99 return &Metrics{ 100 {{- range $metric := .ParsedMetrics }} 101 {{ $metric.FieldName }}: discard.New{{ $metric.TypeName }}(), 102 {{- end }} 103 } 104 } 105 `)) 106 107 // ParsedMetricField is the data parsed for a single field of a metric struct. 108 type ParsedMetricField struct { 109 TypeName string 110 FieldName string 111 MetricName string 112 Description string 113 Labels string 114 115 HistogramOptions HistogramOpts 116 } 117 118 type HistogramOpts struct { 119 BucketType string 120 BucketSizes string 121 } 122 123 // TemplateData is all of the data required for rendering a metric file template. 124 type TemplateData struct { 125 Package string 126 ParsedMetrics []ParsedMetricField 127 } 128 129 func main() { 130 flag.Parse() 131 if *strct == "" { 132 log.Fatal("You must specify a non-empty -struct") 133 } 134 td, err := ParseMetricsDir(".", *strct) 135 if err != nil { 136 log.Fatalf("Parsing file: %v", err) 137 } 138 out := filepath.Join(*dir, "metrics.gen.go") 139 f, err := os.Create(out) 140 if err != nil { 141 log.Fatalf("Opening file: %v", err) 142 } 143 err = GenerateMetricsFile(f, td) 144 if err != nil { 145 log.Fatalf("Generating code: %v", err) 146 } 147 } 148 149 func ignoreTestFiles(f fs.FileInfo) bool { 150 return !strings.Contains(f.Name(), "_test.go") 151 } 152 153 // ParseMetricsDir parses the dir and scans for a struct matching structName, 154 // ignoring all test files. ParseMetricsDir iterates the fields of the metrics 155 // struct and builds a TemplateData using the data obtained from the abstract syntax tree. 156 func ParseMetricsDir(dir string, structName string) (TemplateData, error) { 157 fs := token.NewFileSet() 158 d, err := parser.ParseDir(fs, dir, ignoreTestFiles, parser.ParseComments) 159 if err != nil { 160 return TemplateData{}, err 161 } 162 if len(d) > 1 { 163 return TemplateData{}, fmt.Errorf("multiple packages found in %s", dir) 164 } 165 if len(d) == 0 { 166 return TemplateData{}, fmt.Errorf("no go pacakges found in %s", dir) 167 } 168 169 // Grab the package name. 170 var pkgName string 171 var pkg *ast.Package 172 // TODO(thane): Figure out a more readable way of implementing this. 173 //nolint:revive 174 for pkgName, pkg = range d { 175 } 176 td := TemplateData{ 177 Package: pkgName, 178 } 179 // Grab the metrics struct 180 m, mPkgName, err := findMetricsStruct(pkg.Files, structName) 181 if err != nil { 182 return TemplateData{}, err 183 } 184 for _, f := range m.Fields.List { 185 if !isMetric(f.Type, mPkgName) { 186 continue 187 } 188 pmf := parseMetricField(f) 189 td.ParsedMetrics = append(td.ParsedMetrics, pmf) 190 } 191 192 return td, err 193 } 194 195 // GenerateMetricsFile executes the metrics file template, writing the result 196 // into the io.Writer. 197 func GenerateMetricsFile(w io.Writer, td TemplateData) error { 198 b := []byte{} 199 buf := bytes.NewBuffer(b) 200 err := tmpl.Execute(buf, td) 201 if err != nil { 202 return err 203 } 204 b, err = format.Source(buf.Bytes()) 205 if err != nil { 206 return err 207 } 208 _, err = io.Copy(w, bytes.NewBuffer(b)) 209 if err != nil { 210 return err 211 } 212 return nil 213 } 214 215 func findMetricsStruct(files map[string]*ast.File, structName string) (*ast.StructType, string, error) { 216 var st *ast.StructType 217 for _, file := range files { 218 mPkgName, err := extractMetricsPackageName(file.Imports) 219 if err != nil { 220 return nil, "", fmt.Errorf("unable to determine metrics package name: %v", err) 221 } 222 if !ast.FilterFile(file, func(name string) bool { 223 return name == structName 224 }) { 225 continue 226 } 227 ast.Inspect(file, func(n ast.Node) bool { 228 switch f := n.(type) { 229 case *ast.TypeSpec: 230 if f.Name.Name == structName { 231 var ok bool 232 st, ok = f.Type.(*ast.StructType) 233 if !ok { 234 err = fmt.Errorf("found identifier for %q of wrong type", structName) 235 } 236 } 237 return false 238 default: 239 return true 240 } 241 }) 242 if err != nil { 243 return nil, "", err 244 } 245 if st != nil { 246 return st, mPkgName, nil 247 } 248 } 249 return nil, "", fmt.Errorf("target struct %q not found in dir", structName) 250 } 251 252 func parseMetricField(f *ast.Field) ParsedMetricField { 253 pmf := ParsedMetricField{ 254 Description: extractHelpMessage(f.Doc), 255 MetricName: extractFieldName(f.Names[0].String(), f.Tag), 256 FieldName: f.Names[0].String(), 257 TypeName: extractTypeName(f.Type), 258 Labels: extractLabels(f.Tag), 259 } 260 if pmf.TypeName == "Histogram" { 261 pmf.HistogramOptions = extractHistogramOptions(f.Tag) 262 } 263 return pmf 264 } 265 266 func extractTypeName(e ast.Expr) string { 267 return strings.TrimPrefix(path.Ext(types.ExprString(e)), ".") 268 } 269 270 func extractHelpMessage(cg *ast.CommentGroup) string { 271 if cg == nil { 272 return "" 273 } 274 var help []string //nolint: prealloc 275 for _, c := range cg.List { 276 mt := strings.TrimPrefix(c.Text, "//metrics:") 277 if mt != c.Text { 278 return strings.TrimSpace(mt) 279 } 280 help = append(help, strings.TrimSpace(strings.TrimPrefix(c.Text, "//"))) 281 } 282 return strings.Join(help, " ") 283 } 284 285 func isMetric(e ast.Expr, mPkgName string) bool { 286 return strings.Contains(types.ExprString(e), fmt.Sprintf("%s.", mPkgName)) 287 } 288 289 func extractLabels(bl *ast.BasicLit) string { 290 if bl != nil { 291 t := reflect.StructTag(strings.Trim(bl.Value, "`")) 292 if v := t.Get(labelsTag); v != "" { 293 var res []string 294 for _, s := range strings.Split(v, ",") { 295 res = append(res, strconv.Quote(strings.TrimSpace(s))) 296 } 297 return strings.Join(res, ",") 298 } 299 } 300 return "" 301 } 302 303 func extractFieldName(name string, tag *ast.BasicLit) string { 304 if tag != nil { 305 t := reflect.StructTag(strings.Trim(tag.Value, "`")) 306 if v := t.Get(metricNameTag); v != "" { 307 return v 308 } 309 } 310 return toSnakeCase(name) 311 } 312 313 func extractHistogramOptions(tag *ast.BasicLit) HistogramOpts { 314 h := HistogramOpts{} 315 if tag != nil { 316 t := reflect.StructTag(strings.Trim(tag.Value, "`")) 317 if v := t.Get(bucketTypeTag); v != "" { 318 h.BucketType = bucketType[v] 319 } 320 if v := t.Get(bucketSizeTag); v != "" { 321 h.BucketSizes = v 322 } 323 } 324 return h 325 } 326 327 func extractMetricsPackageName(imports []*ast.ImportSpec) (string, error) { 328 for _, i := range imports { 329 u, err := strconv.Unquote(i.Path.Value) 330 if err != nil { 331 return "", err 332 } 333 if u == metricsPackageName { 334 if i.Name != nil { 335 return i.Name.Name, nil 336 } 337 return path.Base(u), nil 338 } 339 } 340 return "", nil 341 } 342 343 var capitalChange = regexp.MustCompile("([a-z0-9])([A-Z])") 344 345 func toSnakeCase(str string) string { 346 snake := capitalChange.ReplaceAllString(str, "${1}_${2}") 347 return strings.ToLower(snake) 348 }