go/src/infra/cmd/drone-agent/internal/bot/bot.go

// Copyright 2019 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

// Package bot wraps managing Swarming bots.
package bot

import (
	"fmt"
	"io"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"

	"go.chromium.org/luci/common/errors"
)

// Bot is the interface for interacting with a started Swarming bot.
// Wait must be called to ensure the process is waited for.
type Bot interface {
	// Wait waits for the bot process to exit.  The return value
	// on subsequent calls is undefined.
	Wait() error
	// Drain signals for the bot to drain.  Note that this requires
	// support from the bot script.  This should be handled by Swarming
	// bots by waiting for the currently running task to finish before
	// exiting.
	Drain() error
	// Terminate terminates the bot with SIGTERM.  Swarming bots handle
	// SIGTERM by aborting the currently running task and exiting.
	Terminate() error
}

type realBot struct {
	config  Config
	cmd     *exec.Cmd
	logFile *os.File
}

// Wait implements Bot.
func (b realBot) Wait() error {
	err := b.cmd.Wait()
	_ = b.logFile.Close()
	return err
}

// Drain implements Bot.
func (b realBot) Drain() error {
	f, err := os.Create(b.config.drainFilePath())
	if err != nil {
		return errors.Annotate(err, "drain bot %s", b.config.BotID).Err()
	}
	if err := f.Close(); err != nil {
		return errors.Annotate(err, "drain bot %s", b.config.BotID).Err()
	}
	return nil
}

// Starter has a Start method for starting Swarming bots.
type Starter struct {
	client *http.Client
}

// NewStarter returns a new Starter.
func NewStarter(c *http.Client) Starter {
	return Starter{
		client: c,
	}
}

// Start starts a Swarming bot.  The returned Bot object can be used
// to interact with the bot.
func (s Starter) Start(c Config) (b Bot, err error) {
	if err := s.downloadBotCode(c); err != nil {
		return nil, errors.Annotate(err, "start bot with %+v", c).Err()
	}
	f, err := os.Create(c.logFilePath())
	if err != nil {
		return nil, errors.Annotate(err, "start bot with %+v", c).Err()
	}
	defer func() {
		if err != nil {
			_ = f.Close()
		}
	}()
	cmd := exec.Command("python2", c.botZipPath(), "start_bot")
	cmd.Stdout = f
	cmd.Stderr = f
	cmd.Env = append(c.env(), os.Environ()...)
	if err := cmd.Start(); err != nil {
		return nil, errors.Annotate(err, "start bot with %+v", c).Err()
	}
	return realBot{
		config:  c,
		cmd:     cmd,
		logFile: f,
	}, nil
}

func (s Starter) downloadBotCode(c Config) error {
	f, err := os.Create(c.botZipPath())
	if err != nil {
		return errors.Annotate(err, "download bot code for %+v", c).Err()
	}
	defer f.Close()

	resp, err := s.client.Get(c.botCodeURL())
	if err != nil {
		return errors.Annotate(err, "download bot code for %+v", c).Err()
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		return errors.Reason("download bot code for %+v: unexpected status %v", c, resp.StatusCode).Err()
	}

	if _, err := io.Copy(f, resp.Body); err != nil {
		return errors.Annotate(err, "download bot code for %+v", c).Err()
	}
	if err := f.Close(); err != nil {
		return errors.Annotate(err, "download bot code for %+v", c).Err()
	}
	return nil
}

// Config is the configuration needed for starting a generic Swarming bot.
type Config struct {
	// SwarmingURL is the URL of the Swarming instance.  Should be
	// a full URL without the path, e.g. https://host.example.com
	SwarmingURL string
	BotID       string
	// WorkDirectory is the Swarming bot's work directory.
	// The caller should create this.
	// The parent directory should be writable to allow creation
	// of the drain file.
	WorkDirectory string
}

func (c Config) drainFilePath() string {
	return c.WorkDirectory + ".drain"
}

func (c Config) logFilePath() string {
	return c.WorkDirectory + ".log"
}

func (c Config) botZipPath() string {
	return filepath.Join(c.WorkDirectory, "swarming_bot.zip")
}

func (c Config) botCodeURL() string {
	return fmt.Sprintf("%s/bot_code?bot_id=%s", c.SwarmingURL, c.BotID)
}

func (c Config) env() []string {
	return []string{
		"SWARMING_BOT_ID=" + c.BotID,
	}
}