diff --git a/go.mod b/go.mod index df82903f..5e0bdee8 100644 --- a/go.mod +++ b/go.mod @@ -33,7 +33,7 @@ require ( go.opencensus.io v0.0.0-20181129005706-8b019f31bc1c // indirect golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2 golang.org/x/oauth2 v0.0.0-20181128211412-28207608b838 - golang.org/x/sync v0.0.0-20181108010431-42b317875d0f // indirect + golang.org/x/sync v0.0.0-20181108010431-42b317875d0f google.golang.org/api v0.0.0-20181129220737-af4fc4062c26 // indirect google.golang.org/appengine v1.3.0 // indirect google.golang.org/genproto v0.0.0-20181202183823-bd91e49a0898 // indirect diff --git a/infra/gravity/cluster_status.go b/infra/gravity/cluster_status.go index 55c38b81..cd3518d8 100644 --- a/infra/gravity/cluster_status.go +++ b/infra/gravity/cluster_status.go @@ -4,6 +4,7 @@ import ( "context" "github.com/cenkalti/backoff" + "golang.org/x/sync/errgroup" "github.com/gravitational/robotest/lib/constants" sshutils "github.com/gravitational/robotest/lib/ssh" @@ -71,34 +72,32 @@ func (c *TestContext) WaitForStatus(nodes []Gravity, expected statusValidator) e } -// Status queries `gravity status` on each node in nodes. -func (c *TestContext) Status(nodes []Gravity) ([]GravityStatus, error) { +// Status queries `gravity status` once from each node in nodes. +func (c *TestContext) Status(nodes []Gravity) (statuses []GravityStatus, err error) { ctx, cancel := context.WithTimeout(c.ctx, c.timeouts.NodeStatus) defer cancel() - errC := make(chan error, len(nodes)) - valueC := make(chan interface{}, len(nodes)) - + valueC := make(chan GravityStatus, len(nodes)) + g, ctx := errgroup.WithContext(ctx) for _, node := range nodes { - go func(n Gravity) { - status, err := n.Status(ctx) - errC <- err + node := node + g.Go(func() error { + status, err := node.Status(ctx) + if err != nil { + return trace.Wrap(err) + } if status != nil { valueC <- *status } - }(node) + return nil + }) } - - values, err := utils.Collect(ctx, cancel, errC, valueC) + err = g.Wait() if err != nil { return nil, trace.Wrap(err) } - var statuses []GravityStatus - for _, v := range values { - status, ok := v.(GravityStatus) - if !ok { - return nil, trace.BadParameter("expected %T, got %T", status, v) - } + close(valueC) + for status := range valueC { statuses = append(statuses, status) } return statuses, nil