Merge pull request #55 from linuxkit/shard-option

add sharding ability

Author: deitch

cmd/list.go

@@ -31,10 +31,18 @@ var listCmd = &cobra.Command{
 }
 
 func init() {
+	flags := listCmd.Flags()
+	// shardPattern is 1-based (1/10, 3/10, 10/10) rather than normal computer 0-based (0/9, 2/9, 9/9), because it is easier for
+	// humans to understand when calling the CLI.
+	flags.StringVarP(&shardPattern, "shard", "s", "", "which shard to run, in form of 'N/M' where N is the shard number and M is the total number of shards, smallest shard number is 1")
 	RootCmd.AddCommand(listCmd)
 }
 
 func list(_ *cobra.Command, args []string) error {
+	shard, totalShards, err := parseShardPattern(shardPattern)
+	if err != nil {
+		return err
+	}
 	pattern, err := local.ValidatePattern(args)
 	if err != nil {
 		return err
@@ -45,6 +53,11 @@ func list(_ *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	if totalShards > 0 {
+		if err := p.SetShard(shard, totalShards); err != nil {
+			return err
+		}
+	}
 
 	w := new(tabwriter.Writer)
 	w.Init(os.Stdout, 0, 8, 0, '\t', 0)

cmd/run.go

@@ -71,11 +71,12 @@ var (
 )
 
 var (
-	resultDir string
-	id        string
-	symlink   bool
-	extra     bool
-	parallel  bool
+	resultDir    string
+	id           string
+	symlink      bool
+	extra        bool
+	parallel     bool
+	shardPattern string
 )
 
 var runCmd = &cobra.Command{
@@ -90,10 +91,17 @@ func init() {
 	flags.StringVarP(&id, "id", "", "", "ID for this test run")
 	flags.BoolVarP(&extra, "extra", "x", false, "Add extra debug info to log files")
 	flags.BoolVarP(&parallel, "parallel", "p", false, "Run multiple tests in parallel")
+	// shardPattern is 1-based (1/10, 3/10, 10/10) rather than normal computer 0-based (0/9, 2/9, 9/9), because it is easier for
+	// humans to understand when calling the CLI.
+	flags.StringVarP(&shardPattern, "shard", "s", "", "which shard to run, in form of 'N/M' where N is the shard number and M is the total number of shards, smallest shard number is 1")
 	RootCmd.AddCommand(runCmd)
 }
 
 func run(cmd *cobra.Command, args []string) error {
+	shard, totalShards, err := parseShardPattern(shardPattern)
+	if err != nil {
+		return err
+	}
 	pattern, err := local.ValidatePattern(args)
 	if err != nil {
 		return err
@@ -106,6 +114,11 @@ func run(cmd *cobra.Command, args []string) error {
 	if err != nil {
 		return err
 	}
+	if totalShards > 0 {
+		if err := p.SetShard(shard, totalShards); err != nil {
+			return err
+		}
+	}
 
 	var labelList []string
 	for k := range runConfig.Labels {
@@ -309,3 +322,23 @@ func setupResultsDirectory(id string, link bool) (string, error) {
 
 	return baseDir, nil
 }
+
+func parseShardPattern(pattern string) (shard int, total int, err error) {
+	if pattern == "" {
+		return 0, 0, nil
+	}
+	parts := strings.SplitN(pattern, "/", 2)
+	if len(parts) != 2 {
+		return 0, 0, fmt.Errorf("invalid shard pattern: %s", pattern)
+	}
+	if shard, err = strconv.Atoi(parts[0]); err != nil {
+		return 0, 0, fmt.Errorf("invalid shard pattern: %s", pattern)
+	}
+	if total, err = strconv.Atoi(parts[1]); err != nil {
+		return 0, 0, fmt.Errorf("invalid shard pattern: %s", pattern)
+	}
+	if shard < 1 || total < 1 || shard > total {
+		return 0, 0, fmt.Errorf("invalid shard pattern: %s", pattern)
+	}
+	return shard, total, nil
+}

local/group.go

@@ -9,32 +9,8 @@ import (
 	"strings"
 	"sync"
 	"time"
-
-	"github.com/linuxkit/rtf/logger"
 )
 
-// NewProject creates a new top-level Group at the provided path
-func NewProject(path string) (*Group, error) {
-	if !filepath.IsAbs(path) {
-		var err error
-		path, err = filepath.Abs(path)
-		if err != nil {
-			return nil, err
-		}
-	}
-	g := &Group{Parent: nil, Path: path}
-	return g, nil
-}
-
-// InitNewProject creates a new Group, and calls Init() on it
-func InitNewProject(path string) (*Group, error) {
-	group, err := NewProject(path)
-	if err != nil {
-		return group, err
-	}
-	return group, group.Init()
-}
-
 // NewGroup creates a new Group with the given parent and path
 func NewGroup(parent *Group, path string) (*Group, error) {
 	g := &Group{Parent: parent, Path: path, PreTestPath: parent.PreTestPath, PostTestPath: parent.PostTestPath}
@@ -151,44 +127,78 @@ func (g *Group) List(config RunConfig) []Info {
 	return infos
 }
 
+// Gather gathers all runnable child groups and tests
+func (g *Group) Gather(config RunConfig, count int) ([]TestContainer, int) {
+	sort.Sort(ByOrder(g.Children))
+
+	if !g.willRun(config) {
+		return nil, 0
+	}
+	containers := []TestContainer{}
+	var subCount int
+
+	if g.GroupFilePath != "" {
+		containers = append(containers, GroupCommand{Name: g.Name(), FilePath: g.GroupFilePath, Path: g.Path, Type: "init"})
+	}
+
+	for _, c := range g.Children {
+		lst, childCount := c.Gather(config, count+subCount)
+		// if we had no runnable tests, do not bother adding the group init/deinit, just return the empty list
+		if childCount == 0 {
+			continue
+		}
+		containers = append(containers, lst...)
+		subCount += childCount
+	}
+
+	if g.GroupFilePath != "" {
+		containers = append(containers, GroupCommand{Name: g.Name(), FilePath: g.GroupFilePath, Path: g.Path, Type: "deinit"})
+	}
+
+	return containers, subCount
+}
+
 // Run will run all child groups and tests
 func (g *Group) Run(config RunConfig) ([]Result, error) {
 	var results []Result
-	sort.Sort(ByOrder(g.Children))
 
-	if !g.willRun(config) {
+	// This gathers all of the individual tests and group init/deinit commands
+	// all the way down, leading to a flat list we can execute, rather than recursion.
+	// That should make it easier to break into shards.
+	count := 0
+	runnables, _ := g.Gather(config, count)
+	if len(runnables) == 0 {
 		return []Result{{TestResult: Skip,
 			Name:      g.Name(),
 			StartTime: time.Now(),
 			EndTime:   time.Now(),
 		}}, nil
 	}
 
-	if g.GroupFilePath != "" {
-		config.Logger.Log(logger.LevelInfo, fmt.Sprintf("%s::ginit()", g.Name()))
-		res, err := executeScript(g.GroupFilePath, g.Path, "", []string{"init"}, config)
-		if err != nil {
-			return results, err
-		}
-		if res.TestResult != Pass {
-			return results, fmt.Errorf("Error running %s", g.GroupFilePath+":init")
-		}
-	}
-
 	if config.Parallel {
 		var wg sync.WaitGroup
 		resCh := make(chan []Result, len(g.Children))
 		errCh := make(chan error, len(g.Children))
 
-		for _, c := range g.Children {
+		for _, c := range runnables {
 			wg.Add(1)
 			go func(c TestContainer, cf RunConfig) {
 				defer wg.Done()
-				res, err := c.Run(cf)
-				if err != nil {
-					errCh <- err
+				var isTest bool
+				if _, ok := c.(*Test); ok {
+					isTest = true
+				}
+				// Run() if one of the following is true: it is not a test; there are no start/end boundaries; the test is within the boundaries
+				if !isTest || (config.start == 0 && config.count == 0) || (count >= config.start && config.start+config.count > count) {
+					res, err := c.Run(cf)
+					if err != nil {
+						errCh <- err
+					}
+					resCh <- res
+				}
+				if isTest {
+					count++
 				}
-				resCh <- res
 			}(c, config)
 		}
 
@@ -208,25 +218,15 @@ func (g *Group) Run(config RunConfig) ([]Result, error) {
 			results = append(results, res...)
 		}
 	} else {
-		for _, c := range g.Children {
-			res, err := c.Run(config)
+		for _, r := range runnables {
+			res, err := r.Run(config)
 			if err != nil {
 				return results, err
 			}
 			results = append(results, res...)
 		}
 	}
 
-	if g.GroupFilePath != "" {
-		config.Logger.Log(logger.LevelInfo, fmt.Sprintf("%s::gdeinit()", g.Name()))
-		res, err := executeScript(g.GroupFilePath, g.Path, "", []string{"deinit"}, config)
-		if err != nil {
-			return results, err
-		}
-		if res.TestResult != Pass {
-			return results, fmt.Errorf("Error running %s", g.GroupFilePath+":deinit")
-		}
-	}
 	return results, nil
 }
 
@@ -247,3 +247,37 @@ func (g *Group) willRun(config RunConfig) bool {
 
 	return strings.HasPrefix(config.TestPattern, g.Name()) || strings.HasPrefix(g.Name(), config.TestPattern)
 }
+
+func min(x, y int) int {
+	if x < y {
+		return x
+	}
+	return y
+}
+
+// calculateShard calculate the start and end of a slice to use in a given shard of a given total
+// slice and a given number of shards.
+// We split by the total list size. If it is uneven, e.g. 22 elements into 10 shards,
+// then the first shards will be rounded up until the reminder can be split evenly.
+//
+// e.g.
+// 22 elements into 10 shards will be 3, 3, 2, 2, 2, 2, 2, 2, 2, 2
+// 29 elements into 10 shards will be 3, 3, 3, 3, 3, 3, 3, 3, 3, 2
+// 30 elements into 10 shards will be 3, 3, 3, 3, 3, 3, 3, 3, 3, 3
+// 8 elements into 5 shards will be 2, 2, 2, 1, 1
+//
+// Ths important thing is consistency among runs using the same set of parameters
+// so that you can reliably get the same subset each time.
+func calculateShard(size, shard, totalShards int) (start, count int) {
+	elmsPerShard := size / totalShards
+	remainder := size % totalShards
+	before := (shard - 1) * elmsPerShard
+	if remainder > 0 {
+		before += min(remainder, shard-1)
+	}
+	count = elmsPerShard
+	if remainder >= shard {
+		count++
+	}
+	return before, count
+}

local/group_test.go

@@ -0,0 +1,31 @@
+package local
+
+import (
+	"testing"
+)
+
+func TestCalculateShard(t *testing.T) {
+	tests := []struct {
+		size   int
+		shard  int
+		shards int
+		start  int
+		count  int
+	}{
+		{22, 1, 10, 0, 3},
+		{22, 2, 10, 3, 3},
+		{22, 3, 10, 6, 2},
+		{29, 10, 10, 27, 2},
+		{2, 1, 6, 0, 1}, // more shards than elements
+		{2, 2, 6, 1, 1}, // more shards than elements
+		{2, 3, 6, 2, 0}, // more shards than elements
+	}
+	for _, tt := range tests {
+		t.Run("", func(t *testing.T) {
+			if gotStart, gotCount := calculateShard(tt.size, tt.shard, tt.shards); gotStart != tt.start || gotCount != tt.count {
+				t.Errorf("calculateShard() = %v, %v, want %v, %v", gotStart, gotCount, tt.start, tt.count)
+			}
+		})
+	}
+
+}

local/groupcommand.go

@@ -0,0 +1,43 @@
+package local
+
+import (
+	"fmt"
+
+	"github.com/linuxkit/rtf/logger"
+)
+
+// Run satisfies the TestContainer interface.
+// Run the group init or deinit command.
+func (g GroupCommand) Run(config RunConfig) ([]Result, error) {
+	config.Logger.Log(logger.LevelInfo, fmt.Sprintf("%s::%s()", g.Name, g.Type))
+	res, err := executeScript(g.FilePath, g.Path, "", []string{g.Type}, config)
+	if err != nil {
+		return nil, err
+	}
+	if res.TestResult != Pass {
+		return nil, fmt.Errorf("error running %s:%s", g.FilePath, g.Type)
+	}
+	return []Result{res}, nil
+}
+
+// List satisfies the TestContainer interface
+func (g GroupCommand) List(config RunConfig) []Info {
+	info := Info{
+		Name: g.Name,
+	}
+
+	return []Info{info}
+}
+
+// Order returns a tests order
+func (g GroupCommand) Order() int {
+	if g.Type == "init" {
+		return 0
+	}
+	return 1
+}
+
+// Gather satisfies the TestContainer interface
+func (g GroupCommand) Gather(config RunConfig, count int) ([]TestContainer, int) {
+	return []TestContainer{&g}, 0
+}