agent-tools/waitfor/wait/wait.go

// Package wait blocks until a shell condition holds: run a command
// every interval until it exits 0 (and, optionally, its output matches
// a regex) or a timeout expires. One blocking call instead of an
// agent burning turns on poll loops.
package wait

import (
	"fmt"
	"os/exec"
	"regexp"
	"time"
)

type Options struct {
	Cmd      string        // shell command to run each attempt (required)
	Matches  string        // optional regex the output must match
	Interval time.Duration // between attempts
	Timeout  time.Duration // total budget
	Verbose  bool          // progress line per attempt to the log func
}

type Result struct {
	Met      bool
	Attempts int
	Elapsed  time.Duration
	LastOut  string
	LastExit int
}

// Runner executes a shell command, returning combined output and exit
// code. Separated out so tests can fake it.
type Runner func(cmd string) (string, int)

// ShellRunner runs via sh -c with combined stdout+stderr.
func ShellRunner(cmd string) (string, int) {
	c := exec.Command("sh", "-c", cmd)
	out, err := c.CombinedOutput()
	code := 0
	if err != nil {
		code = 1
		if ee, ok := err.(*exec.ExitError); ok {
			code = ee.ExitCode()
		}
	}
	return string(out), code
}

// Wait polls until the condition holds or the timeout expires. The
// first attempt runs immediately. log receives progress lines when
// Verbose is set (pass nil otherwise).
func Wait(opts Options, run Runner, log func(string)) (Result, error) {
	if opts.Cmd == "" {
		return Result{}, fmt.Errorf("no command given")
	}
	var re *regexp.Regexp
	if opts.Matches != "" {
		var err error
		re, err = regexp.Compile(opts.Matches)
		if err != nil {
			return Result{}, fmt.Errorf("bad --matches regex: %v", err)
		}
	}
	if opts.Interval <= 0 {
		opts.Interval = 10 * time.Second
	}
	if opts.Timeout <= 0 {
		opts.Timeout = 10 * time.Minute
	}
	start := time.Now()
	res := Result{}
	for {
		res.Attempts++
		out, code := run(opts.Cmd)
		res.LastOut, res.LastExit = out, code
		met := code == 0 && (re == nil || re.MatchString(out))
		if opts.Verbose && log != nil {
			state := "not yet"
			if met {
				state = "met"
			}
			log(fmt.Sprintf("attempt %d: exit %d, condition %s", res.Attempts, code, state))
		}
		if met {
			res.Met = true
			res.Elapsed = time.Since(start)
			return res, nil
		}
		if time.Since(start)+opts.Interval > opts.Timeout {
			res.Elapsed = time.Since(start)
			return res, nil
		}
		time.Sleep(opts.Interval)
	}
}