Add sharding support to multiproduct_kati

This is so that we can split the build_test build on large branches into
multiple builds, each testing a subset of the products.

Test: m blueprint_tools  (runs the go unit tests)
Test: build/soong/build_test.bash --only-config
      vs
Test: build/soong/build_test.bash --only-config --shard-count=2
Test: build/soong/build_test.bash --only-config --shard-count=2 --shard=2

Change-Id: I40ccc1aa477bc0ffa74ff564d155068509be18f0

3 files changed, 125 insertions, 0 deletions

diff --git a/cmd/multiproduct_kati/Android.bp b/cmd/multiproduct_kati/Android.bp
index 13b3679e..d34f8c3c 100644
--- a/cmd/multiproduct_kati/Android.bp
+++ b/cmd/multiproduct_kati/Android.bp
@@ -24,4 +24,7 @@ blueprint_go_binary {
     srcs: [
         "main.go",
     ],
+    testSrcs: [
+        "main_test.go",
+    ],
 }
diff --git a/cmd/multiproduct_kati/main.go b/cmd/multiproduct_kati/main.go
index 2800ade2..47712062 100644
--- a/cmd/multiproduct_kati/main.go
+++ b/cmd/multiproduct_kati/main.go
@@ -64,6 +64,9 @@ var buildVariant = flag.String("variant", "eng", "build variant to use")
 var skipProducts = flag.String("skip-products", "", "comma-separated list of products to skip (known failures, etc)")
 var includeProducts = flag.String("products", "", "comma-separated list of products to build")
 
+var shardCount = flag.Int("shard-count", 1, "split the products into multiple shards (to spread the build onto multiple machines, etc)")
+var shard = flag.Int("shard", 1, "1-indexed shard to execute")
+
 const errorLeadingLines = 20
 const errorTrailingLines = 20
 
@@ -278,6 +281,17 @@ func main() {
 		}
 	}
 
+	if *shard < 1 {
+		log.Fatalf("--shard value must be >= 1, not %d\n", *shard)
+	} else if *shardCount < 1 {
+		log.Fatalf("--shard-count value must be >= 1, not %d\n", *shardCount)
+	} else if *shard > *shardCount {
+		log.Fatalf("--shard (%d) must not be greater than --shard-count (%d)\n", *shard,
+			*shardCount)
+	} else if *shardCount > 1 {
+		finalProductsList = splitList(finalProductsList, *shardCount)[*shard-1]
+	}
+
 	log.Verbose("Got product list: ", finalProductsList)
 
 	s := buildCtx.Status.StartTool()
@@ -472,3 +486,18 @@ func (f *failureCount) Write(p []byte) (int, error) {
 	// discard writes
 	return len(p), nil
 }
+
+func splitList(list []string, shardCount int) (ret [][]string) {
+	each := len(list) / shardCount
+	extra := len(list) % shardCount
+	for i := 0; i < shardCount; i++ {
+		count := each
+		if extra > 0 {
+			count += 1
+			extra -= 1
+		}
+		ret = append(ret, list[:count])
+		list = list[count:]
+	}
+	return
+}
diff --git a/cmd/multiproduct_kati/main_test.go b/cmd/multiproduct_kati/main_test.go
new file mode 100644
index 00000000..263a124f
--- /dev/null
+++ b/cmd/multiproduct_kati/main_test.go
@@ -0,0 +1,93 @@
+// Copyright 2019 Google Inc. All rights reserved.
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package main
+
+import (
+	"fmt"
+	"reflect"
+	"testing"
+)
+
+func TestSplitList(t *testing.T) {
+	testcases := []struct {
+		inputCount int
+		shardCount int
+		want       [][]string
+	}{
+		{
+			inputCount: 1,
+			shardCount: 1,
+			want:       [][]string{{"1"}},
+		},
+		{
+			inputCount: 1,
+			shardCount: 2,
+			want:       [][]string{{"1"}, {}},
+		},
+		{
+			inputCount: 4,
+			shardCount: 2,
+			want:       [][]string{{"1", "2"}, {"3", "4"}},
+		},
+		{
+			inputCount: 19,
+			shardCount: 10,
+			want: [][]string{
+				{"1", "2"},
+				{"3", "4"},
+				{"5", "6"},
+				{"7", "8"},
+				{"9", "10"},
+				{"11", "12"},
+				{"13", "14"},
+				{"15", "16"},
+				{"17", "18"},
+				{"19"},
+			},
+		},
+		{
+			inputCount: 15,
+			shardCount: 10,
+			want: [][]string{
+				{"1", "2"},
+				{"3", "4"},
+				{"5", "6"},
+				{"7", "8"},
+				{"9", "10"},
+				{"11"},
+				{"12"},
+				{"13"},
+				{"14"},
+				{"15"},
+			},
+		},
+	}
+
+	for _, tc := range testcases {
+		t.Run(fmt.Sprintf("%d/%d", tc.inputCount, tc.shardCount), func(t *testing.T) {
+			input := []string{}
+			for i := 1; i <= tc.inputCount; i++ {
+				input = append(input, fmt.Sprintf("%d", i))
+			}
+
+			got := splitList(input, tc.shardCount)
+
+			if !reflect.DeepEqual(got, tc.want) {
+				t.Errorf("unexpected result for splitList([]string{...%d...}, %d):\nwant: %v\n got: %v\n",
+					tc.inputCount, tc.shardCount, tc.want, got)
+			}
+		})
+	}
+}