Add --invoke option to launch a container from the build result

Signed-off-by: Kohei Tokunaga <ktokunaga.mail@gmail.com>
This commit is contained in:
Kohei Tokunaga
2022-05-24 20:00:55 +09:00
parent 1343cdfc83
commit 387e1ecca6
5 changed files with 1007 additions and 8 deletions

486
monitor/monitor.go Normal file
View File

@ -0,0 +1,486 @@
package monitor
import (
"bufio"
"context"
"fmt"
"io"
"sync"
"github.com/docker/buildx/build"
"github.com/pkg/errors"
"github.com/sirupsen/logrus"
"golang.org/x/term"
)
// RunMonitor provides an interactive session for running and managing containers via specified IO.
func RunMonitor(ctx context.Context, containerConfig build.ContainerConfig, reloadFunc func(context.Context) (*build.ResultContext, error), stdin io.ReadCloser, stdout, stderr io.WriteCloser) error {
monitorIn, monitorOut := ioSetPipe()
defer monitorIn.Close()
monitorEnableCh := make(chan struct{})
monitorDisableCh := make(chan struct{})
monitorOutCtx := ioSetOutContext{monitorOut,
func() { monitorEnableCh <- struct{}{} },
func() { monitorDisableCh <- struct{}{} },
}
containerIn, containerOut := ioSetPipe()
defer containerIn.Close()
containerOutCtx := ioSetOutContext{containerOut,
// send newline to hopefully get the prompt; TODO: better UI (e.g. reprinting the last line)
func() { containerOut.stdin.Write([]byte("\n")) },
func() {},
}
m := &monitor{
invokeIO: newIOForwarder(containerIn),
muxIO: newMuxIO(ioSetIn{stdin, stdout, stderr}, []ioSetOutContext{monitorOutCtx, containerOutCtx}, 1, "Switched IO\n"),
}
// Start container automatically
go func() {
m.rollback(ctx, containerConfig)
}()
// Serve monitor commands
monitorForwarder := newIOForwarder(monitorIn)
for {
<-monitorEnableCh
in, out := ioSetPipe()
monitorForwarder.setDestination(&out)
doneCh, errCh := make(chan struct{}), make(chan error)
go func() {
defer close(doneCh)
defer in.Close()
t := term.NewTerminal(readWriter{in.stdin, in.stdout}, "(buildx) ")
for {
l, err := t.ReadLine()
if err != nil {
if err != io.EOF {
errCh <- err
return
}
return
}
switch l {
case "":
// nop
case "reload":
res, err := reloadFunc(ctx)
if err != nil {
fmt.Printf("failed to reload: %v\n", err)
} else {
// rollback the running container with the new result
containerConfig.ResultCtx = res
m.rollback(ctx, containerConfig)
}
case "rollback":
m.rollback(ctx, containerConfig)
case "exit":
return
default:
fmt.Printf("unknown command: %q\n", l)
}
}
}()
select {
case <-doneCh:
return nil
case err := <-errCh:
return err
case <-monitorDisableCh:
}
monitorForwarder.setDestination(nil)
}
}
type readWriter struct {
io.Reader
io.Writer
}
type monitor struct {
muxIO *muxIO
invokeIO *ioForwarder
curInvokeCancel func()
}
func (m *monitor) rollback(ctx context.Context, cfg build.ContainerConfig) {
if m.curInvokeCancel != nil {
m.curInvokeCancel() // Finish the running container if exists
}
go func() {
// Start a new container
if err := m.invoke(ctx, cfg); err != nil {
logrus.Debugf("invoke error: %v", err)
}
}()
}
func (m *monitor) invoke(ctx context.Context, cfg build.ContainerConfig) error {
m.muxIO.enable(1)
defer m.muxIO.disable(1)
invokeCtx, invokeCancel := context.WithCancel(ctx)
containerIn, containerOut := ioSetPipe()
m.invokeIO.setDestination(&containerOut)
waitInvokeDoneCh := make(chan struct{})
var cancelOnce sync.Once
curInvokeCancel := func() {
cancelOnce.Do(func() {
containerIn.Close()
m.invokeIO.setDestination(nil)
invokeCancel()
})
<-waitInvokeDoneCh
}
defer curInvokeCancel()
m.curInvokeCancel = curInvokeCancel
cfg.Stdin = containerIn.stdin
cfg.Stdout = containerIn.stdout
cfg.Stderr = containerIn.stderr
err := build.Invoke(invokeCtx, cfg)
close(waitInvokeDoneCh)
return err
}
type ioForwarder struct {
curIO *ioSetOut
mu sync.Mutex
updateCh chan struct{}
}
func newIOForwarder(in ioSetIn) *ioForwarder {
f := &ioForwarder{
updateCh: make(chan struct{}),
}
doneCh := make(chan struct{})
go func() {
for {
f.mu.Lock()
w := f.curIO
f.mu.Unlock()
if w != nil && w.stdout != nil && w.stderr != nil {
go func() {
if _, err := io.Copy(in.stdout, w.stdout); err != nil && err != io.ErrClosedPipe {
// ErrClosedPipe is OK as we close this read end during setDestination.
logrus.WithError(err).Warnf("failed to forward stdout: %v", err)
}
}()
go func() {
if _, err := io.Copy(in.stderr, w.stderr); err != nil && err != io.ErrClosedPipe {
// ErrClosedPipe is OK as we close this read end during setDestination.
logrus.WithError(err).Warnf("failed to forward stderr: %v", err)
}
}()
}
select {
case <-f.updateCh:
case <-doneCh:
return
}
}
}()
go func() {
if err := copyToFunc(in.stdin, func() (io.Writer, error) {
f.mu.Lock()
w := f.curIO
f.mu.Unlock()
if w != nil {
return w.stdin, nil
}
return nil, nil
}); err != nil && err != io.ErrClosedPipe {
logrus.WithError(err).Warnf("failed to forward IO: %v", err)
}
close(doneCh)
if w := f.curIO; w != nil {
// Propagate close
if err := w.Close(); err != nil {
logrus.WithError(err).Warnf("failed to forwarded stdin IO: %v", err)
}
}
}()
return f
}
func (f *ioForwarder) setDestination(out *ioSetOut) {
f.mu.Lock()
if f.curIO != nil {
// close all stream on the current IO no to mix with the new IO
f.curIO.Close()
}
f.curIO = out
f.mu.Unlock()
f.updateCh <- struct{}{}
}
type ioSetOutContext struct {
ioSetOut
enableHook func()
disableHook func()
}
// newMuxIO forwards IO stream to/from "in" and "outs".
// "outs" are closed automatically when "in" reaches EOF.
// "in" doesn't closed automatically so the caller needs to explicitly close it.
func newMuxIO(in ioSetIn, out []ioSetOutContext, initIdx int, toggleMessage string) *muxIO {
m := &muxIO{
enabled: make(map[int]struct{}),
in: in,
out: out,
closedCh: make(chan struct{}),
toggleMessage: toggleMessage,
}
for i := range out {
m.enabled[i] = struct{}{}
}
m.maxCur = len(out)
m.cur = initIdx
var wg sync.WaitGroup
var mu sync.Mutex
for i, o := range out {
i, o := i, o
wg.Add(1)
go func() {
defer wg.Done()
if err := copyToFunc(o.stdout, func() (io.Writer, error) {
if m.cur == i {
return in.stdout, nil
}
return nil, nil
}); err != nil {
logrus.WithField("output index", i).WithError(err).Warnf("failed to write stdout")
}
if err := o.stdout.Close(); err != nil {
logrus.WithField("output index", i).WithError(err).Warnf("failed to close stdout")
}
}()
wg.Add(1)
go func() {
defer wg.Done()
if err := copyToFunc(o.stderr, func() (io.Writer, error) {
if m.cur == i {
return in.stderr, nil
}
return nil, nil
}); err != nil {
logrus.WithField("output index", i).WithError(err).Warnf("failed to write stderr")
}
if err := o.stderr.Close(); err != nil {
logrus.WithField("output index", i).WithError(err).Warnf("failed to close stderr")
}
}()
}
go func() {
errToggle := errors.Errorf("toggle IO")
for {
prevIsControlSequence := false
if err := copyToFunc(traceReader(in.stdin, func(r rune) (bool, error) {
// Toggle IO when it detects C-a-c
// TODO: make it configurable if needed
if int(r) == 1 {
prevIsControlSequence = true
return false, nil
}
defer func() { prevIsControlSequence = false }()
if prevIsControlSequence {
if string(r) == "c" {
return false, errToggle
}
}
return true, nil
}), func() (io.Writer, error) {
mu.Lock()
o := out[m.cur]
mu.Unlock()
return o.stdin, nil
}); !errors.Is(err, errToggle) {
if err != nil {
logrus.WithError(err).Warnf("failed to read stdin")
}
break
}
m.toggleIO()
}
// propagate stdin EOF
for i, o := range out {
if err := o.stdin.Close(); err != nil {
logrus.WithError(err).Warnf("failed to close stdin of %d", i)
}
}
wg.Wait()
close(m.closedCh)
}()
return m
}
type muxIO struct {
cur int
maxCur int
enabled map[int]struct{}
mu sync.Mutex
in ioSetIn
out []ioSetOutContext
closedCh chan struct{}
toggleMessage string
}
func (m *muxIO) waitClosed() {
<-m.closedCh
}
func (m *muxIO) enable(i int) {
m.mu.Lock()
defer m.mu.Unlock()
m.enabled[i] = struct{}{}
}
func (m *muxIO) disable(i int) error {
m.mu.Lock()
defer m.mu.Unlock()
if i == 0 {
return errors.Errorf("disabling 0th io is prohibited")
}
delete(m.enabled, i)
if m.cur == i {
m.toggleIO()
}
return nil
}
func (m *muxIO) toggleIO() {
if m.out[m.cur].disableHook != nil {
m.out[m.cur].disableHook()
}
for {
if m.cur+1 >= m.maxCur {
m.cur = 0
} else {
m.cur++
}
if _, ok := m.enabled[m.cur]; !ok {
continue
}
break
}
if m.out[m.cur].enableHook != nil {
m.out[m.cur].enableHook()
}
fmt.Fprintf(m.in.stdout, m.toggleMessage)
}
func traceReader(r io.ReadCloser, f func(rune) (bool, error)) io.ReadCloser {
pr, pw := io.Pipe()
go func() {
br := bufio.NewReader(r)
for {
rn, _, err := br.ReadRune()
if err != nil {
if err == io.EOF {
pw.Close()
return
}
pw.CloseWithError(err)
return
}
if isWrite, err := f(rn); err != nil {
pw.CloseWithError(err)
return
} else if !isWrite {
continue
}
if _, err := pw.Write([]byte(string(rn))); err != nil {
pw.CloseWithError(err)
return
}
}
}()
return &readerWithClose{
Reader: pr,
closeFunc: func() error {
pr.Close()
return r.Close()
},
}
}
func copyToFunc(r io.Reader, wFunc func() (io.Writer, error)) error {
buf := make([]byte, 4096)
for {
n, readErr := r.Read(buf)
if readErr != nil && readErr != io.EOF {
return readErr
}
w, err := wFunc()
if err != nil {
return err
}
if w != nil {
if _, err := w.Write(buf[:n]); err != nil {
logrus.WithError(err).Debugf("failed to copy")
}
}
if readErr == io.EOF {
return nil
}
}
}
func ioSetPipe() (ioSetIn, ioSetOut) {
r1, w1 := io.Pipe()
r2, w2 := io.Pipe()
r3, w3 := io.Pipe()
return ioSetIn{r1, w2, w3}, ioSetOut{w1, r2, r3}
}
type ioSetIn struct {
stdin io.ReadCloser
stdout io.WriteCloser
stderr io.WriteCloser
}
func (s ioSetIn) Close() (retErr error) {
if err := s.stdin.Close(); err != nil {
retErr = err
}
if err := s.stdout.Close(); err != nil {
retErr = err
}
if err := s.stderr.Close(); err != nil {
retErr = err
}
return
}
type ioSetOut struct {
stdin io.WriteCloser
stdout io.ReadCloser
stderr io.ReadCloser
}
func (s ioSetOut) Close() (retErr error) {
if err := s.stdin.Close(); err != nil {
retErr = err
}
if err := s.stdout.Close(); err != nil {
retErr = err
}
if err := s.stderr.Close(); err != nil {
retErr = err
}
return
}
type readerWithClose struct {
io.Reader
closeFunc func() error
}
func (r *readerWithClose) Close() error {
return r.closeFunc()
}

321
monitor/monitor_test.go Normal file
View File

@ -0,0 +1,321 @@
package monitor
import (
"bytes"
"fmt"
"io"
"regexp"
"strings"
"testing"
"golang.org/x/sync/errgroup"
)
// TestMuxIO tests muxIO
func TestMuxIO(t *testing.T) {
tests := []struct {
name string
inputs []instruction
initIdx int
outputsNum int
wants []string
// Everytime string is written to the mux stdin, the output end
// that received the string write backs to the string that is masked with
// its index number. This is useful to check if writeback is written from the
// expected output destination.
wantsMaskedOutput string
}{
{
name: "single output",
inputs: []instruction{
input("foo\nbar\n"),
toggle(),
input("1234"),
toggle(),
input("456"),
},
initIdx: 0,
outputsNum: 1,
wants: []string{"foo\nbar\n1234456"},
wantsMaskedOutput: `^0+$`,
},
{
name: "multi output",
inputs: []instruction{
input("foo\nbar\n"),
toggle(),
input("12" + string([]rune{rune(1)}) + "34abc"),
toggle(),
input("456"),
},
initIdx: 0,
outputsNum: 3,
wants: []string{"foo\nbar\n", "1234abc", "456"},
wantsMaskedOutput: `^0+1+2+$`,
},
{
name: "multi output with nonzero index",
inputs: []instruction{
input("foo\nbar\n"),
toggle(),
input("1234"),
toggle(),
input("456"),
},
initIdx: 1,
outputsNum: 3,
wants: []string{"456", "foo\nbar\n", "1234"},
wantsMaskedOutput: `^1+2+0+$`,
},
{
name: "multi output many toggles",
inputs: []instruction{
input("foo\nbar\n"),
toggle(),
input("1234"),
toggle(),
toggle(),
input("456"),
toggle(),
input("%%%%"),
toggle(),
toggle(),
toggle(),
input("aaaa"),
},
initIdx: 0,
outputsNum: 3,
wants: []string{"foo\nbar\n456", "1234%%%%aaaa", ""},
wantsMaskedOutput: `^0+1+0+1+$`,
},
{
name: "enable disable",
inputs: []instruction{
input("foo\nbar\n"),
toggle(),
input("1234"),
toggle(),
input("456"),
disable(2),
input("%%%%"),
enable(2),
toggle(),
toggle(),
input("aaa"),
disable(2),
disable(1),
input("1111"),
toggle(),
input("2222"),
toggle(),
input("3333"),
},
initIdx: 0,
outputsNum: 3,
wants: []string{"foo\nbar\n%%%%111122223333", "1234", "456aaa"},
wantsMaskedOutput: `^0+1+2+0+2+0+$`,
},
}
for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {
inBuf, end, in := newTestIn(t)
var outBufs []*outBuf
var outs []ioSetOutContext
if tt.outputsNum != len(tt.wants) {
t.Fatalf("wants != outputsNum")
}
for i := 0; i < tt.outputsNum; i++ {
outBuf, out := newTestOut(t, i)
outBufs = append(outBufs, outBuf)
outs = append(outs, ioSetOutContext{out, nil, nil})
}
mio := newMuxIO(in, outs, tt.initIdx, "")
for _, i := range tt.inputs {
// Add input to muxIO
istr, writeback := i(mio)
if _, err := end.stdin.Write([]byte(istr)); err != nil {
t.Fatalf("failed to write data to stdin: %v", err)
}
// Wait for writeback of this input
var eg errgroup.Group
eg.Go(func() error {
outbuf := make([]byte, len(writeback))
if _, err := io.ReadAtLeast(end.stdout, outbuf, len(outbuf)); err != nil {
return err
}
return nil
})
eg.Go(func() error {
errbuf := make([]byte, len(writeback))
if _, err := io.ReadAtLeast(end.stderr, errbuf, len(errbuf)); err != nil {
return err
}
return nil
})
if err := eg.Wait(); err != nil {
t.Fatalf("failed to wait for output: %v", err)
}
}
// Close stdin on this muxIO
end.stdin.Close()
// Wait for all output ends reach EOF
mio.waitClosed()
// Close stdout/stderr as well
in.Close()
// Check if each output end received expected string
<-inBuf.doneCh
for i, o := range outBufs {
<-o.doneCh
if o.stdin != tt.wants[i] {
t.Fatalf("output[%d]: got %q; wanted %q", i, o.stdin, tt.wants[i])
}
}
// Check if expected string is returned from expected outputs
if !regexp.MustCompile(tt.wantsMaskedOutput).MatchString(inBuf.stdout) {
t.Fatalf("stdout: got %q; wanted %q", inBuf.stdout, tt.wantsMaskedOutput)
}
if !regexp.MustCompile(tt.wantsMaskedOutput).MatchString(inBuf.stderr) {
t.Fatalf("stderr: got %q; wanted %q", inBuf.stderr, tt.wantsMaskedOutput)
}
})
}
}
type instruction func(m *muxIO) (intput string, writeBackView string)
func input(s string) instruction {
return func(m *muxIO) (string, string) {
return s, strings.ReplaceAll(s, string([]rune{rune(1)}), "")
}
}
func toggle() instruction {
return func(m *muxIO) (string, string) {
return string([]rune{rune(1)}) + "c", ""
}
}
func enable(i int) instruction {
return func(m *muxIO) (string, string) {
m.enable(i)
return "", ""
}
}
func disable(i int) instruction {
return func(m *muxIO) (string, string) {
m.disable(i)
return "", ""
}
}
type inBuf struct {
stdout string
stderr string
doneCh chan struct{}
}
func newTestIn(t *testing.T) (*inBuf, ioSetOut, ioSetIn) {
ti := &inBuf{
doneCh: make(chan struct{}),
}
gotOutR, gotOutW := io.Pipe()
gotErrR, gotErrW := io.Pipe()
outR, outW := io.Pipe()
var eg errgroup.Group
eg.Go(func() error {
buf := new(bytes.Buffer)
if _, err := io.Copy(io.MultiWriter(gotOutW, buf), outR); err != nil {
return err
}
ti.stdout = buf.String()
return nil
})
errR, errW := io.Pipe()
eg.Go(func() error {
buf := new(bytes.Buffer)
if _, err := io.Copy(io.MultiWriter(gotErrW, buf), errR); err != nil {
return err
}
ti.stderr = buf.String()
return nil
})
go func() {
eg.Wait()
close(ti.doneCh)
}()
inR, inW := io.Pipe()
return ti, ioSetOut{inW, gotOutR, gotErrR}, ioSetIn{inR, outW, errW}
}
type outBuf struct {
idx int
stdin string
doneCh chan struct{}
}
func newTestOut(t *testing.T, idx int) (*outBuf, ioSetOut) {
to := &outBuf{
idx: idx,
doneCh: make(chan struct{}),
}
inR, inW := io.Pipe()
outR, outW := io.Pipe()
errR, errW := io.Pipe()
go func() {
defer inR.Close()
defer outW.Close()
defer errW.Close()
buf := new(bytes.Buffer)
mw := io.MultiWriter(buf,
writeMasked(outW, fmt.Sprintf("%d", to.idx)),
writeMasked(errW, fmt.Sprintf("%d", to.idx)),
)
if _, err := io.Copy(mw, inR); err != nil {
inR.CloseWithError(err)
outW.CloseWithError(err)
errW.CloseWithError(err)
return
}
to.stdin = string(buf.Bytes())
outW.Close()
errW.Close()
close(to.doneCh)
}()
return to, ioSetOut{inW, outR, errR}
}
func writeMasked(w io.Writer, s string) io.Writer {
buf := make([]byte, 4096)
pr, pw := io.Pipe()
go func() {
for {
n, readErr := pr.Read(buf)
if readErr != nil && readErr != io.EOF {
pr.CloseWithError(readErr)
return
}
var masked string
for i := 0; i < n; i++ {
masked += s
}
if _, err := w.Write([]byte(masked)); err != nil {
pr.CloseWithError(err)
return
}
if readErr == io.EOF {
pr.Close()
return
}
}
}()
return pw
}