intake/core/execute.go

186 lines
4.4 KiB
Go

package core
import (
"bufio"
"context"
"encoding/json"
"errors"
"fmt"
"io"
"log"
"os"
"os/exec"
"strings"
"time"
)
func readStdout(stdout io.ReadCloser, source string, items chan Item, cparse chan bool) {
var item Item
parseError := false
scanout := bufio.NewScanner(stdout)
for scanout.Scan() {
data := scanout.Bytes()
err := json.Unmarshal(data, &item)
if err != nil || item.Id == "" {
log.Printf("[%s: stdout] %s\n", source, strings.TrimSpace(string(data)))
parseError = true
} else {
item.Active = true // These fields aren't up to
item.Created = 0 // the action to set and
item.Source = source // shouldn't be overrideable
log.Printf("[%s: item] %s\n", source, item.Id)
items <- item
}
}
// Only send the parsing result at the end, to block main until stdout is drained
cparse <- parseError
close(items)
}
func readStderr(stderr io.ReadCloser, source string, done chan bool) {
scanerr := bufio.NewScanner(stderr)
for scanerr.Scan() {
text := strings.TrimSpace(scanerr.Text())
log.Printf("[%s: stderr] %s\n", source, text)
}
done <- true
}
func writeStdin(stdin io.WriteCloser, text string) {
defer stdin.Close()
io.WriteString(stdin, text)
}
func Execute(
source string,
argv []string,
env []string,
state []byte,
input string,
timeout time.Duration,
) ([]Item, []byte, error) {
log.Printf("executing %v", argv)
if len(argv) == 0 {
return nil, nil, errors.New("empty argv")
}
if source == "" {
return nil, nil, errors.New("empty source")
}
stateFile, err := os.CreateTemp("", "intake_state_*")
if err != nil {
return nil, nil, fmt.Errorf("error: failed to create temp state file: %v", err)
}
defer func() {
if err := os.Remove(stateFile.Name()); err != nil {
log.Printf("error: failed to delete %s", stateFile.Name())
}
}()
_, err = stateFile.Write(state)
if err != nil {
return nil, nil, fmt.Errorf("error: failed to write state file: %v", err)
}
env = append(env, "STATE_PATH="+stateFile.Name())
ctx, cancel := context.WithTimeout(context.Background(), timeout)
defer cancel()
cmd := exec.CommandContext(ctx, argv[0], argv[1:]...)
cmd.Env = append(os.Environ(), env...)
cmd.WaitDelay = time.Second * 5
// Open pipes to the command
stdin, err := cmd.StdinPipe()
if err != nil {
return nil, nil, err
}
stdout, err := cmd.StdoutPipe()
if err != nil {
return nil, nil, err
}
stderr, err := cmd.StderrPipe()
if err != nil {
return nil, nil, err
}
cout := make(chan Item)
cparse := make(chan bool)
cerr := make(chan bool)
// Sink routine for items produced
var items []Item
go func() {
for item := range cout {
items = append(items, item)
}
}()
// Routines handling the process i/o
go writeStdin(stdin, input)
go readStdout(stdout, source, cout, cparse)
go readStderr(stderr, source, cerr)
// Kick off the command
err = cmd.Start()
if err != nil {
return nil, nil, err
}
// Block until std{out,err} close
<-cerr
parseError := <-cparse
err = cmd.Wait()
if ctx.Err() == context.DeadlineExceeded {
log.Printf("Timed out after %v\n", timeout)
return nil, nil, err
} else if exiterr, ok := err.(*exec.ExitError); ok {
log.Printf("error: %s failed with exit code %d\n", argv[0], exiterr.ExitCode())
return nil, nil, err
} else if err != nil {
log.Printf("error: %s failed with error: %s\n", argv[0], err)
return nil, nil, err
}
if parseError {
log.Printf("error: could not parse item\n")
return nil, nil, errors.New("invalid JSON")
}
newState, err := os.ReadFile(stateFile.Name())
if err != nil {
return nil, nil, fmt.Errorf("error: failed to read state file: %v", err)
}
return items, newState, nil
}
// Execute an action that takes an item as input and returns the item modified.
// This is basically just a wrapper over [Execute] that handles the input and backfilling.
func ExecuteItemAction(
item Item,
argv []string,
env []string,
state []byte,
timeout time.Duration,
) (Item, []byte, error) {
itemJson, err := json.Marshal(item)
if err != nil {
return Item{}, nil, fmt.Errorf("failed to serialize item: %v", err)
}
res, newState, err := Execute(item.Source, argv, env, state, string(itemJson), timeout)
if err != nil {
return Item{}, nil, fmt.Errorf("failed to execute action for %s/%s: %v", item.Source, item.Id, err)
}
if len(res) != 1 {
return Item{}, nil, fmt.Errorf("expected action to produce exactly one item, got %d", len(res))
}
newItem := res[0]
BackfillItem(&newItem, &item)
return newItem, newState, nil
}