task_driver/examples/file_stream/file_stream.go

package main

/*
	FileStream example.

	Run like this:

	$ go run ./file_stream.go --logtostderr --project_id=skia-swarming-bots --task_name=filestream_example -o - --local
*/

import (
	"context"
	"flag"
	"fmt"
	"io"
	"os"
	"path/filepath"
	"runtime"

	"go.skia.org/infra/go/exec"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
	"go.skia.org/infra/task_driver/go/td"
)

var (
	// Required flags for all TaskDrivers.
	projectId = flag.String("project_id", "", "ID of the Google Cloud project.")
	taskId    = flag.String("task_id", "", "ID of this task.")
	taskName  = flag.String("task_name", "", "Name of the task.")
	output    = flag.String("o", "", "If provided, dump a JSON blob of step data to the given file. Prints to stdout if '-' is given.")
	local     = flag.Bool("local", false, "True if running locally (as opposed to in production)")

	// Flags for this TaskDriver.
	workdir = flag.String("workdir", os.TempDir(), "Working directory.")
)

func main() {
	// Setup.
	taskName := "FileStream Example"
	ctx := td.StartRun(projectId, taskId, &taskName, output, local)
	defer td.EndRun(ctx)

	if err := example1(ctx); err != nil {
		td.Fatal(ctx, err)
	}
	if err := example2(ctx); err != nil {
		td.Fatal(ctx, err)
	}
}

func example1(ctx context.Context) (rv error) {
	ctx = td.StartStep(ctx, td.Props("example1"))
	defer td.EndStep(ctx)

	// This script writes logs to a file.
	_, filename, _, ok := runtime.Caller(1)
	if !ok {
		return td.FailStep(ctx, fmt.Errorf("Failed to obtain path of current file."))
	}
	script := filepath.Join(filepath.Dir(filename), "write_logs.py")
	fs, err := td.NewFileStream(ctx, "verbose", sklog.DEBUG)
	if err != nil {
		return td.FailStep(ctx, err)
	}
	defer util.Close(fs)
	if _, err := exec.RunCwd(ctx, *workdir, "python", "-u", script, fs.FilePath()); err != nil {
		return td.FailStep(ctx, err)
	}
	return nil
}

func example2(ctx context.Context) (rv error) {
	ctx = td.StartStep(ctx, td.Props("example2"))
	defer td.EndStep(ctx)

	// File streams should also work when the file is copied over.
	tmpFile := filepath.Join(*workdir, "tmpfile")
	if err := util.WithWriteFile(tmpFile, func(w io.Writer) error {
		if _, err := w.Write([]byte("Contents were copied (via os.Rename)")); err != nil {
			return td.FailStep(ctx, err)
		}
		return nil
	}); err != nil {
		return td.FailStep(ctx, err)
	}
	fs, err := td.NewFileStream(ctx, "copied", sklog.DEBUG)
	if err != nil {
		return td.FailStep(ctx, err)
	}
	defer util.Close(fs)
	if _, err := exec.RunCwd(ctx, *workdir, "cp", tmpFile, fs.FilePath()); err != nil {
		return td.FailStep(ctx, err)
	}
	return nil
}