go.chromium.org/luci@v0.0.0-20240309015107-7cdc2e660f33/gce/cmd/agent/swarming.go (about)

     1  // Copyright 2018 The LUCI Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //      http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    12  // See the License for the specific language governing permissions and
    13  // limitations under the License.
    14  
    15  package main
    16  
    17  import (
    18  	"context"
    19  	"io"
    20  	"net/http"
    21  	"os"
    22  	"path/filepath"
    23  
    24  	"go.chromium.org/luci/common/errors"
    25  	"go.chromium.org/luci/common/logging"
    26  )
    27  
    28  // SwarmingClient is a Swarming server client.
    29  type SwarmingClient struct {
    30  	// Client is the *http.Client to use to communicate with the Swarming server.
    31  	*http.Client
    32  	// PlatformStrategy is the platform-specific strategy to use.
    33  	PlatformStrategy
    34  	// server is the Swarming server URL.
    35  	server string
    36  }
    37  
    38  // swrKey is the key to a *SwarmingClient in the context.
    39  var swrKey = "swr"
    40  
    41  // withSwarming returns a new context with the given *SwarmingClient installed.
    42  func withSwarming(c context.Context, cli *SwarmingClient) context.Context {
    43  	return context.WithValue(c, &swrKey, cli)
    44  }
    45  
    46  // getSwarming returns the *SwarmingClient installed in the current context.
    47  func getSwarming(c context.Context) *SwarmingClient {
    48  	return c.Value(&swrKey).(*SwarmingClient)
    49  }
    50  
    51  // fetch fetches the Swarming bot code.
    52  func (s *SwarmingClient) fetch(c context.Context, path, user string) error {
    53  	botCode := s.server + "/bot_code"
    54  	logging.Infof(c, "downloading: %s", botCode)
    55  	rsp, err := s.Get(botCode)
    56  	if err != nil {
    57  		return errors.Annotate(err, "failed to fetch bot code").Err()
    58  	}
    59  	defer rsp.Body.Close()
    60  	if rsp.StatusCode != http.StatusOK {
    61  		return errors.Reason("server returned %q", rsp.Status).Err()
    62  	}
    63  
    64  	logging.Infof(c, "installing: %s", path)
    65  	// 0644 allows the bot code to be read by all users.
    66  	// Useful when SSHing to the instance.
    67  	out, err := os.OpenFile(path, os.O_CREATE|os.O_WRONLY, 0644)
    68  	if err != nil {
    69  		return errors.Annotate(err, "failed to open: %s", path).Err()
    70  	}
    71  	defer out.Close()
    72  	_, err = io.Copy(out, rsp.Body)
    73  	if err != nil {
    74  		return errors.Annotate(err, "failed to write: %s", path).Err()
    75  	}
    76  	if err := s.chown(c, path, user); err != nil {
    77  		return errors.Annotate(err, "failed to chown: %s", path).Err()
    78  	}
    79  	return nil
    80  }
    81  
    82  // Configure fetches the Swarming bot code and configures it to run on startup.
    83  func (s *SwarmingClient) Configure(c context.Context, dir, user string, python string) error {
    84  	// 0755 allows the directory structure to be read and listed by all users.
    85  	// Useful when SSHing fo the instance.
    86  	if err := os.MkdirAll(dir, 0755); err != nil {
    87  		return errors.Annotate(err, "failed to create: %s", dir).Err()
    88  	}
    89  	if err := s.chown(c, dir, user); err != nil {
    90  		return errors.Annotate(err, "failed to chown: %s", dir).Err()
    91  	}
    92  	zip := filepath.Join(dir, "swarming_bot.zip")
    93  	switch _, err := os.Stat(zip); {
    94  	case os.IsNotExist(err):
    95  	case err != nil:
    96  		return errors.Annotate(err, "failed to stat: %s", zip).Err()
    97  	default:
    98  		logging.Infof(c, "already installed: %s", zip)
    99  		return nil
   100  	}
   101  	if err := s.fetch(c, zip, user); err != nil {
   102  		return err
   103  	}
   104  	return s.autostart(c, zip, user, python)
   105  }