Add --ninja_weight_source option

It has 3 options for now.
* ninja_log: uses ninja log file for data source
* evenly_distributed: pass empty list for ninja to consider every work as the same
* not_used: do not use critical path logic in ninja

In addition, I added the option in the metrics to track

Bug: 271527305
Test: m --ninja_weight_source=ninja_log|empty|not_used
Change-Id: Ib4c812c20606a34b17d3f0edb71057b477c4f90e
diff --git a/ui/build/config.go b/ui/build/config.go
index d8a2788..8a77b08 100644
--- a/ui/build/config.go
+++ b/ui/build/config.go
@@ -118,8 +118,21 @@
 	bazelForceEnabledModules string
 
 	includeTags []string
+
+	// Data source to write ninja weight list
+	ninjaWeightListSource NinjaWeightListSource
 }
 
+type NinjaWeightListSource uint
+
+const (
+	// ninja doesn't use weight list.
+	NOT_USED NinjaWeightListSource = iota
+	// ninja uses weight list based on previous builds by ninja log
+	NINJA_LOG
+	// ninja thinks every task has the same weight.
+	EVENLY_DISTRIBUTED
+)
 const srcDirFileCheck = "build/soong/root.bp"
 
 var buildFiles = []string{"Android.mk", "Android.bp"}
@@ -528,6 +541,17 @@
 	ctx.Metrics.SystemResourceInfo(s)
 }
 
+func getNinjaWeightListSourceInMetric(s NinjaWeightListSource) *smpb.NinjaWeightListSource {
+	switch s {
+	case NINJA_LOG:
+		return smpb.NinjaWeightListSource_NINJA_LOG.Enum()
+	case EVENLY_DISTRIBUTED:
+		return smpb.NinjaWeightListSource_EVENLY_DISTRIBUTED.Enum()
+	default:
+		return smpb.NinjaWeightListSource_NOT_USED.Enum()
+	}
+}
+
 func buildConfig(config Config) *smpb.BuildConfig {
 	c := &smpb.BuildConfig{
 		ForceUseGoma:                proto.Bool(config.ForceUseGoma()),
@@ -535,6 +559,7 @@
 		UseRbe:                      proto.Bool(config.UseRBE()),
 		BazelMixedBuild:             proto.Bool(config.BazelBuildEnabled()),
 		ForceDisableBazelMixedBuild: proto.Bool(config.IsBazelMixedBuildForceDisabled()),
+		NinjaWeightListSource:       getNinjaWeightListSourceInMetric(config.NinjaWeightListSource()),
 	}
 	c.Targets = append(c.Targets, config.arguments...)
 
@@ -797,6 +822,17 @@
 			c.bazelStagingMode = true
 		} else if arg == "--search-api-dir" {
 			c.searchApiDir = true
+		} else if strings.HasPrefix(arg, "--ninja_weight_source=") {
+			source := strings.TrimPrefix(arg, "--ninja_weight_source=")
+			if source == "ninja_log" {
+				c.ninjaWeightListSource = NINJA_LOG
+			} else if source == "evenly_distributed" {
+				c.ninjaWeightListSource = EVENLY_DISTRIBUTED
+			} else if source == "not_used" {
+				c.ninjaWeightListSource = NOT_USED
+			} else {
+				ctx.Fatalf("unknown option for ninja_weight_source: %s", source)
+			}
 		} else if strings.HasPrefix(arg, "--build-command=") {
 			buildCmd := strings.TrimPrefix(arg, "--build-command=")
 			// remove quotations
@@ -1091,6 +1127,10 @@
 	return c.multitreeBuild
 }
 
+func (c *configImpl) NinjaWeightListSource() NinjaWeightListSource {
+	return c.ninjaWeightListSource
+}
+
 func (c *configImpl) SkipKati() bool {
 	return c.skipKati
 }
diff --git a/ui/build/config_test.go b/ui/build/config_test.go
index 1749cd4..bedb181 100644
--- a/ui/build/config_test.go
+++ b/ui/build/config_test.go
@@ -1032,6 +1032,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(false),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1043,6 +1044,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(false),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1054,6 +1056,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(false),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1065,6 +1068,7 @@
 				UseRbe:                      proto.Bool(true),
 				BazelMixedBuild:             proto.Bool(false),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1076,6 +1080,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(false),
 				ForceDisableBazelMixedBuild: proto.Bool(true),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1088,6 +1093,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(false),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1100,6 +1106,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(true),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1112,6 +1119,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(true),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1124,6 +1132,7 @@
 				UseRbe:                      proto.Bool(false),
 				BazelMixedBuild:             proto.Bool(true),
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1138,6 +1147,7 @@
 				BazelMixedBuild:             proto.Bool(false),
 				Targets:                     []string{"droid", "dist"},
 				ForceDisableBazelMixedBuild: proto.Bool(false),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 		{
@@ -1156,6 +1166,7 @@
 				UseRbe:                      proto.Bool(true),
 				BazelMixedBuild:             proto.Bool(true),
 				ForceDisableBazelMixedBuild: proto.Bool(true),
+				NinjaWeightListSource:       smpb.NinjaWeightListSource_NOT_USED.Enum(),
 			},
 		},
 	}
diff --git a/ui/build/ninja.go b/ui/build/ninja.go
index 28f3c38..a91cc3b 100644
--- a/ui/build/ninja.go
+++ b/ui/build/ninja.go
@@ -33,6 +33,46 @@
 	ninjaEnvFileName = "ninja.environment"
 )
 
+func useNinjaBuildLog(ctx Context, config Config, cmd *Cmd) {
+	ninjaLogFile := filepath.Join(config.OutDir(), ".ninja_log")
+	data, err := os.ReadFile(ninjaLogFile)
+	var outputBuilder strings.Builder
+	if err == nil {
+		lines := strings.Split(strings.TrimSpace(string(data)), "\n")
+		// ninja log: <start>	<end>	<restat>	<name>	<cmdhash>
+		// ninja weight list: <name>,<end-start+1>
+		for _, line := range lines {
+			if strings.HasPrefix(line, "#") {
+				continue
+			}
+			fields := strings.Split(line, "\t")
+			path := fields[3]
+			start, err := strconv.Atoi(fields[0])
+			if err != nil {
+				continue
+			}
+			end, err := strconv.Atoi(fields[1])
+			if err != nil {
+				continue
+			}
+			outputBuilder.WriteString(path)
+			outputBuilder.WriteString(",")
+			outputBuilder.WriteString(strconv.Itoa(end-start+1) + "\n")
+		}
+	}
+	// If there is no ninja log file, just pass empty ninja weight list.
+	// Because it is still efficient with critical path calculation logic even without weight.
+
+	weightListFile := filepath.Join(config.OutDir(), ".ninja_weight_list")
+
+	err = os.WriteFile(weightListFile, []byte(outputBuilder.String()), 0644)
+	if err == nil {
+		cmd.Args = append(cmd.Args, "-o", "usesweightlist="+weightListFile)
+	} else {
+		ctx.Panicf("Could not write ninja weight list file %s", err)
+	}
+}
+
 // Constructs and runs the Ninja command line with a restricted set of
 // environment variables. It's important to restrict the environment Ninja runs
 // for hermeticity reasons, and to avoid spurious rebuilds.
@@ -85,6 +125,14 @@
 		cmd.Environment.AppendFromKati(config.KatiEnvFile())
 	}
 
+	switch config.NinjaWeightListSource() {
+	case NINJA_LOG:
+		useNinjaBuildLog(ctx, config, cmd)
+	case EVENLY_DISTRIBUTED:
+		// pass empty weight list means ninja considers every tasks's weight as 1(default value).
+		cmd.Args = append(cmd.Args, "-o", "usesweightlist=/dev/null")
+	}
+
 	// Allow both NINJA_ARGS and NINJA_EXTRA_ARGS, since both have been
 	// used in the past to specify extra ninja arguments.
 	if extra, ok := cmd.Environment.Get("NINJA_ARGS"); ok {