Update the prepare_prebuilt_workers to be able to specify repo.

tools/README.md

@@ -98,13 +98,27 @@ bin/prepare_prebuilt_workers \
      -r containers/pre_built_workers
 ```
 
-This builds `cxx` and `go` images contain workers built from the commit/branch
-we wish to test.
+This builds `cxx` and `go` images contain workers built from the specific
+commit/branch we wish to test.
+
+You may also specify the repo where the source code is pulled from by specifying
+the repo. The following command pulls the source code from forked repos
+`${USER}/grpc` and `${USER}/grpc-go` to build cxx and go images
+
+```shell
+bin/prepare_prebuilt_workers \
+     -l cxx:"${USER}"/grpc:master \
+     -l go:"${USER}"/grpc-go:master \
+     -p "${image_registry}" \
+     -t "${tag}" \
+     -r containers/pre_built_workers
+```
 
 The tool `prepare_prebuilt_workers` takes the following options:
 
 - `-l`<br> Language and GITREF to benchmark. The language and its specific
-  GITREF wish to build workers from can be specified as `language:COMMIT_SHA`.
+  GITREF wish to build workers from can be specified as `language:COMMIT_SHA` or
+  `language:repo:COMMIT_SHA`.
   May be repeated. Valid input for language names are all in lower case:
   `csharp`, `c++`/`cxx`, `go`, `java`, `node`, `node_purejs`, `php7`.
   `php7_protobuf_c`, `python`, `python_asyncio` and `ruby`.

tools/cmd/prepare_prebuilt_workers/main.go

@@ -37,11 +37,18 @@ import (
 // Tests contains the values for fields that are accessible by
 // flags.
 type Tests struct {
-	preBuiltImagePrefix string
-	testTag             string
-	dockerfileRoot      string
-	buildOnly           bool
-	languagesToGitref   map[string]string
+	preBuiltImagePrefix     string
+	testTag                 string
+	dockerfileRoot          string
+	buildOnly               bool
+	languagesToLanguageSpec map[string]LanguageSpec
+}
+
+// LanguageSpec containers the specs of each tested language.
+type LanguageSpec struct {
+	Name   string `json:"name"`
+	Repo   string `json:"repo"`
+	Gitref string `json:"gitref"`
 }
 
 type langFlags []string
@@ -72,7 +79,7 @@ func main() {
 
 	flag.StringVar(&test.dockerfileRoot, "r", "", "root directory of Dockerfiles to build prebuilt images")
 
-	flag.Var(&languagesSelected, "l", "languages and its GITREF wish to run tests, example: cxx:<commit-sha>")
+	flag.Var(&languagesSelected, "l", "languages, its repository and GITREF wish to run tests, example: cxx:<commit-sha> or cxx:grpc/grpc:<commit-sha>")
 
 	flag.Parse()
 
@@ -94,52 +101,63 @@ func main() {
 		log.Fatalf("Failed preparing prebuilt images: no language and its gitref pair specified, please provide languages and the GITREF as cxx:master")
 	}
 
-	test.languagesToGitref = map[string]string{}
+	test.languagesToLanguageSpec = map[string]LanguageSpec{}
 	converterToImageLanguage := map[string]string{
 		"c++":             "cxx",
 		"node_purejs":     "node",
 		"php7_protobuf_c": "php7",
 		"python_asyncio":  "python",
 	}
 
-	for _, pair := range languagesSelected {
-		split := strings.Split(pair, ":")
+	for _, s := range languagesSelected {
+		split := strings.SplitN(s, ":", 3)
 
-		if len(split) != 2 || split[len(split)-1] == "" {
-			log.Fatalf("Input error in language and gitref selection, please follow the format as language:gitref, for example: c++:<commit-sha>")
+		// C++:master will be split to 2 items, c++:grpc/grpc:master will be
+		// split to 3 items.
+		if len(split) < 2 || split[len(split)-1] == "" {
+			log.Fatalf("Input error in language and gitref selection. Please follow the format language:gitref or language:repository:gitref, for example c++:master or c++:grpc/grpc:master")
 		}
-
+		var spec LanguageSpec
 		lang := split[0]
-		gitref := split[1]
-
 		if convertedLang, ok := converterToImageLanguage[lang]; ok {
-			test.languagesToGitref[convertedLang] = gitref
+			spec.Name = convertedLang
 		} else {
-			test.languagesToGitref[lang] = gitref
+			spec.Name = lang
 		}
+		if len(split) == 3 {
+			spec.Repo = split[1]
+			spec.Gitref = split[2]
+		} else {
+			spec.Repo = ""
+			spec.Gitref = split[1]
+		}
+		test.languagesToLanguageSpec[spec.Name] = spec
 	}
 
-	log.Println("Selected language : GITREF")
-	formattedMap, _ := json.MarshalIndent(test.languagesToGitref, "", "  ")
+	log.Println("Selected language : REPOSITORY: GITREF")
+	formattedMap, _ := json.MarshalIndent(test.languagesToLanguageSpec, "", "  ")
 	log.Print(string(formattedMap))
 
 	var wg sync.WaitGroup
-	wg.Add(len(test.languagesToGitref))
+	wg.Add(len(test.languagesToLanguageSpec))
 
 	uniqueCacheBreaker := time.Now().String()
 
-	for lang, gitRef := range test.languagesToGitref {
-		go func(lang string, gitRef string) {
+	for lang, spec := range test.languagesToLanguageSpec {
+		go func(lang string, spec LanguageSpec) {
 			defer wg.Done()
 
 			image := fmt.Sprintf("%s/%s:%s", test.preBuiltImagePrefix, lang, test.testTag)
 			dockerfileLocation := fmt.Sprintf("%s/%s/", test.dockerfileRoot, lang)
 
-			// build image
-			log.Println(fmt.Sprintf("building %s image", lang))
-			// TODO(jtattermusch): allow setting REPOSITORY arg to allow building from user forks e.g. https://github.com/USER_NAME/grpc.git
+			// Build image
+			log.Printf("building %s image\n", lang)
 			buildCommandTimeoutSeconds := 30 * 60 // 30 mins should be enough for all languages
-			buildDockerImage := exec.Command("timeout", fmt.Sprintf("%ds", buildCommandTimeoutSeconds), "docker", "build", dockerfileLocation, "-t", image, "--build-arg", fmt.Sprintf("GITREF=%s", gitRef), "--build-arg", fmt.Sprintf("BREAK_CACHE=%s", uniqueCacheBreaker))
+			buildDockerImage := exec.Command("timeout", fmt.Sprintf("%ds", buildCommandTimeoutSeconds), "docker", "build", dockerfileLocation, "-t", image, "--build-arg", fmt.Sprintf("GITREF=%s", spec.Gitref), "--build-arg", fmt.Sprintf("BREAK_CACHE=%s", uniqueCacheBreaker))
+			if spec.Repo != "" {
+				buildDockerImage.Args = append(buildDockerImage.Args, "--build-arg", fmt.Sprintf("REPOSITORY=%s", spec.Repo))
+			}
+			log.Printf("Running command: %s", strings.Join(buildDockerImage.Args, " "))
 			buildOutput, err := buildDockerImage.CombinedOutput()
 			if err != nil {
 				log.Printf("Failed building %s image. Dump of command's output will follow:\n", lang)
@@ -151,8 +169,8 @@ func main() {
 			log.Printf("Succeeded building %s image: %s\n", lang, image)
 
 			if !test.buildOnly {
-				// push image
-				log.Println(fmt.Sprintf("pushing %s image", lang))
+				// Push image
+				log.Printf("pushing %s image\n", lang)
 				pushDockerImage := exec.Command("docker", "push", image)
 				pushOutput, err := pushDockerImage.CombinedOutput()
 				if err != nil {
@@ -164,7 +182,7 @@ func main() {
 				log.Println(string(pushOutput))
 				log.Printf("Succeeded pushing %s image to %s\n", lang, image)
 			}
-		}(lang, gitRef)
+		}(lang, spec)
 	}
 
 	wg.Wait()