golang.org/x/build@v0.0.0-20240506185731-218518f32b70/internal/task/swarming.go (about)

     1  // Copyright 2023 The Go Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  package task
     6  
     7  import (
     8  	"context"
     9  	"fmt"
    10  	"strings"
    11  
    12  	"go.chromium.org/luci/swarming/client/swarming"
    13  	apipb "go.chromium.org/luci/swarming/proto/api_v2"
    14  )
    15  
    16  type SwarmingClient interface {
    17  	// RunTask runs script on a machine running port with env set.
    18  	// The script will have the latest version of Go and some version of gsutil
    19  	// on $PATH. To facilitate Windows/Unix compatibility, . will be at the end
    20  	// of $PATH.
    21  	RunTask(ctx context.Context, dims map[string]string, script string, env map[string]string) (string, error)
    22  	// Completed reports whether a build has finished, returning an error if
    23  	// it's failed. It's suitable for use with AwaitCondition.
    24  	Completed(ctx context.Context, id string) (string, bool, error)
    25  }
    26  
    27  type RealSwarmingClient struct {
    28  	SwarmingClient                           swarming.Client
    29  	SwarmingURL, ServiceAccount, Realm, Pool string
    30  }
    31  
    32  func (c *RealSwarmingClient) RunTask(ctx context.Context, dims map[string]string, script string, env map[string]string) (string, error) {
    33  	cipdPlatform, ok := dims["cipd_platform"]
    34  	if !ok {
    35  		return "", fmt.Errorf("must specify cipd_platform in dims: %v", dims)
    36  	}
    37  	shell := []string{"bash", "-eux"}
    38  	if strings.HasPrefix(cipdPlatform, "windows") {
    39  		shell = []string{"cmd", "/c"}
    40  	}
    41  
    42  	req := &apipb.NewTaskRequest{
    43  		Name:           "relui task",
    44  		ExpirationSecs: 3 * 60 * 60,
    45  		Priority:       20,
    46  		User:           "relui",
    47  		ServiceAccount: c.ServiceAccount,
    48  		Realm:          c.Realm,
    49  
    50  		Properties: &apipb.TaskProperties{
    51  			EnvPrefixes: []*apipb.StringListPair{
    52  				{Key: "PATH", Value: []string{"tools/bin"}},
    53  			},
    54  			Env:     []*apipb.StringPair{},
    55  			Command: append(append([]string{"luci-auth", "context"}, shell...), script),
    56  			CipdInput: &apipb.CipdInput{
    57  				Packages: []*apipb.CipdPackage{
    58  					{Path: "tools/bin", PackageName: "infra/tools/luci-auth/" + cipdPlatform, Version: "latest"},
    59  					{Path: "tools", PackageName: "golang/bootstrap-go/" + cipdPlatform, Version: "latest"},
    60  					{Path: "tools", PackageName: "infra/3pp/tools/gcloud/" + cipdPlatform, Version: "latest"},
    61  					{Path: "tools", PackageName: "infra/3pp/tools/cpython3/" + cipdPlatform, Version: "latest"},
    62  				},
    63  			},
    64  			Dimensions: []*apipb.StringPair{
    65  				{Key: "pool", Value: c.Pool},
    66  			},
    67  			ExecutionTimeoutSecs: 600,
    68  		},
    69  	}
    70  	for k, v := range dims {
    71  		req.Properties.Dimensions = append(req.Properties.Dimensions, &apipb.StringPair{Key: k, Value: v})
    72  	}
    73  	for k, v := range env {
    74  		req.Properties.Env = append(req.Properties.Env, &apipb.StringPair{Key: k, Value: v})
    75  	}
    76  	task, err := c.SwarmingClient.NewTask(ctx, req)
    77  	if err != nil {
    78  		return "", err
    79  	}
    80  	return task.TaskId, nil
    81  }
    82  
    83  func (c *RealSwarmingClient) Completed(ctx context.Context, id string) (string, bool, error) {
    84  	result, err := c.SwarmingClient.TaskResult(ctx, id, &swarming.TaskResultFields{WithPerf: false})
    85  	if err != nil {
    86  		return "", false, err
    87  	}
    88  	if result.State == apipb.TaskState_RUNNING || result.State == apipb.TaskState_PENDING {
    89  		return "", false, nil
    90  	}
    91  	if result.State != apipb.TaskState_COMPLETED || result.ExitCode != 0 {
    92  		return "", true, fmt.Errorf("build failed with state %v and exit code %v, see %v/task?id=%v", result.State, result.ExitCode, c.SwarmingURL, id)
    93  	}
    94  	return "", true, nil
    95  }