Carlos Florêncio
127 lines
3.4 KiB
Go
127 lines
3.4 KiB
Go
package process
|
|
|
|
import (
|
|
"bytes"
|
|
"context"
|
|
"errors"
|
|
"io"
|
|
"time"
|
|
|
|
"github.com/kandev/kandev/internal/agentctl/types"
|
|
tools "github.com/kandev/kandev/internal/tools/installer"
|
|
)
|
|
|
|
const managedCommandOutputLimit = 1 << 20
|
|
|
|
// CommandEnvironment returns the task environment overrides applied to every
|
|
// command owned by this manager.
|
|
func (m *Manager) CommandEnvironment() (map[string]string, error) {
|
|
return mergeAgentEnvIntoShellConfigWithError(m.agentEnvSnapshot(), nil)
|
|
}
|
|
|
|
// Output runs a direct command as an instance-owned process and returns bounded
|
|
// stdout after its complete process tree is reaped. Stderr is drained but is
|
|
// not joined to stdout, so callers that parse stdout cannot mistake a warning
|
|
// path for command output.
|
|
func (m *Manager) Output(parent context.Context, spec tools.CommandSpec) ([]byte, error) {
|
|
stdout, _, err := m.runManagedCommand(parent, spec)
|
|
return stdout, err
|
|
}
|
|
|
|
// CombinedOutput runs a direct command as an instance-owned process and
|
|
// returns bounded stdout/stderr after its complete process tree is reaped.
|
|
func (m *Manager) CombinedOutput(parent context.Context, spec tools.CommandSpec) ([]byte, error) {
|
|
stdout, stderr, err := m.runManagedCommand(parent, spec)
|
|
return append(stdout, stderr...), err
|
|
}
|
|
|
|
func (m *Manager) runManagedCommand(parent context.Context, spec tools.CommandSpec) ([]byte, []byte, error) {
|
|
ctx, release, err := m.BeginOwnedOperation(parent)
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
defer release()
|
|
|
|
sessionID := m.cfg.SessionID
|
|
if sessionID == "" {
|
|
sessionID = m.cfg.InstanceID
|
|
}
|
|
if sessionID == "" {
|
|
sessionID = "managed-command"
|
|
}
|
|
if err := ctx.Err(); err != nil {
|
|
return nil, nil, err
|
|
}
|
|
proc, err := m.StartPipedProcess(PipedStartRequest{
|
|
SessionID: sessionID,
|
|
Kind: types.ProcessKindCustom,
|
|
ScriptName: "managed-command",
|
|
Command: spec.Path,
|
|
Args: spec.Args,
|
|
WorkingDir: spec.Dir,
|
|
Env: spec.Env,
|
|
PipeStderr: true,
|
|
})
|
|
if err != nil {
|
|
return nil, nil, err
|
|
}
|
|
_ = proc.Stdin.Close()
|
|
stdout := collectManagedCommandOutput(proc.Stdout)
|
|
stderr := collectManagedCommandOutput(proc.Stderr)
|
|
|
|
var stopErr error
|
|
canceled := false
|
|
finished := true
|
|
select {
|
|
case <-proc.Done:
|
|
case <-ctx.Done():
|
|
canceled = true
|
|
stopCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
|
|
stopErr = m.StopProcess(stopCtx, StopProcessRequest{ProcessID: proc.ID})
|
|
select {
|
|
case <-proc.Done:
|
|
case <-stopCtx.Done():
|
|
finished = false
|
|
stopErr = errors.Join(stopErr, stopCtx.Err())
|
|
}
|
|
cancel()
|
|
}
|
|
if !finished {
|
|
return nil, nil, errors.Join(ctx.Err(), stopErr)
|
|
}
|
|
|
|
stdoutOutput := <-stdout
|
|
stderrOutput := <-stderr
|
|
waitErr := proc.Wait()
|
|
if canceled {
|
|
return stdoutOutput, stderrOutput, errors.Join(ctx.Err(), stopErr, waitErr)
|
|
}
|
|
return stdoutOutput, stderrOutput, waitErr
|
|
}
|
|
|
|
func collectManagedCommandOutput(reader io.ReadCloser) <-chan []byte {
|
|
result := make(chan []byte, 1)
|
|
go func() {
|
|
defer func() { _ = reader.Close() }()
|
|
var output bytes.Buffer
|
|
_, _ = io.Copy(&boundedOutputWriter{buffer: &output, remaining: managedCommandOutputLimit}, reader)
|
|
result <- output.Bytes()
|
|
}()
|
|
return result
|
|
}
|
|
|
|
type boundedOutputWriter struct {
|
|
buffer *bytes.Buffer
|
|
remaining int
|
|
}
|
|
|
|
func (w *boundedOutputWriter) Write(data []byte) (int, error) {
|
|
written := len(data)
|
|
if w.remaining > 0 {
|
|
keep := min(len(data), w.remaining)
|
|
_, _ = w.buffer.Write(data[:keep])
|
|
w.remaining -= keep
|
|
}
|
|
return written, nil
|
|
}
|