diff --git a/engine/remote.go b/engine/remote.go index 60aad243c5..fbb9951065 100644 --- a/engine/remote.go +++ b/engine/remote.go @@ -90,19 +90,21 @@ func (rcv *Receiver) Run() error { f.Close() return err } + f.Close() + defer peer.Close() cmd := data.Message(p).Get("cmd") job := rcv.Engine.Job(cmd[0], cmd[1:]...) - stdout, err := beam.SendPipe(peer, data.Empty().Set("cmd", "log", "stdout").Bytes()) + stdout, err := beam.SendRPipe(peer, data.Empty().Set("cmd", "log", "stdout").Bytes()) if err != nil { return err } job.Stdout.Add(stdout) - stderr, err := beam.SendPipe(peer, data.Empty().Set("cmd", "log", "stderr").Bytes()) + stderr, err := beam.SendRPipe(peer, data.Empty().Set("cmd", "log", "stderr").Bytes()) if err != nil { return err } job.Stderr.Add(stderr) - stdin, err := beam.SendPipe(peer, data.Empty().Set("cmd", "log", "stdin").Bytes()) + stdin, err := beam.SendWPipe(peer, data.Empty().Set("cmd", "log", "stdin").Bytes()) if err != nil { return err } diff --git a/engine/remote_test.go b/engine/remote_test.go index 54092ec934..921a4376ce 100644 --- a/engine/remote_test.go +++ b/engine/remote_test.go @@ -1,3 +1,94 @@ package engine -import () +import ( + "bufio" + "bytes" + "fmt" + "github.com/dotcloud/docker/pkg/beam" + "strings" + "testing" + "time" +) + +func TestHelloWorld(t *testing.T) { + for i := 0; i < 10; i++ { + testRemote(t, + + // Sender side + func(eng *Engine) { + job := eng.Job("echo", "hello", "world") + out := &bytes.Buffer{} + job.Stdout.Add(out) + job.Run() + if job.status != StatusOK { + t.Fatalf("#%v", job.StatusCode()) + } + lines := bufio.NewScanner(out) + var i int + for lines.Scan() { + if lines.Text() != "hello world" { + t.Fatalf("%#v", lines.Text()) + } + i++ + } + if i != 1000 { + t.Fatalf("%#v", i) + } + }, + + // Receiver side + func(eng *Engine) { + eng.Register("echo", func(job *Job) Status { + // Simulate more output with a delay in the middle + for i := 0; i < 500; i++ { + fmt.Fprintf(job.Stdout, "%s\n", strings.Join(job.Args, " ")) + } + time.Sleep(5 * time.Millisecond) + for i := 0; i < 500; i++ { + fmt.Fprintf(job.Stdout, "%s\n", strings.Join(job.Args, " ")) + } + return StatusOK + }) + }, + ) + } +} + +// Helpers + +func testRemote(t *testing.T, senderSide, receiverSide func(*Engine)) { + sndConn, rcvConn, err := beam.USocketPair() + if err != nil { + t.Fatal(err) + } + defer sndConn.Close() + defer rcvConn.Close() + sender := NewSender(sndConn) + receiver := NewReceiver(rcvConn) + + // Setup the sender side + eng := New() + sender.Install(eng) + + // Setup the receiver side + receiverSide(receiver.Engine) + go receiver.Run() + + timeout(t, func() { + senderSide(eng) + }) +} + +func timeout(t *testing.T, f func()) { + onTimeout := time.After(100 * time.Millisecond) + onDone := make(chan bool) + go func() { + f() + close(onDone) + }() + select { + case <-onTimeout: + t.Fatalf("timeout") + case <-onDone: + } +} diff --git a/pkg/beam/beam.go b/pkg/beam/beam.go index b1e4667a3f..88d3ea7c79 100644 --- a/pkg/beam/beam.go +++ b/pkg/beam/beam.go @@ -29,17 +29,48 @@ type ReceiveSender interface { Sender } -func SendPipe(dst Sender, data []byte) (*os.File, error) { +const ( + R int = 1 << (32 - 1 - iota) + W +) + +func sendPipe(dst Sender, data []byte, mode int) (*os.File, error) { r, w, err := os.Pipe() if err != nil { return nil, err } - if err := dst.Send(data, r); err != nil { - r.Close() - w.Close() + var ( + remote *os.File + local *os.File + ) + if mode == R { + remote = r + local = w + } else if mode == W { + remote = w + local = r + } + if err := dst.Send(data, remote); err != nil { + local.Close() + remote.Close() return nil, err } - return w, nil + return local, nil + +} + +// SendRPipe create a pipe and sends its *read* end attached in a beam message +// to `dst`, with `data` as the message payload. +// It returns the *write* end of the pipe, or an error. +func SendRPipe(dst Sender, data []byte) (*os.File, error) { + return sendPipe(dst, data, R) +} + +// SendWPipe create a pipe and sends its *read* end attached in a beam message +// to `dst`, with `data` as the message payload. +// It returns the *write* end of the pipe, or an error. +func SendWPipe(dst Sender, data []byte) (*os.File, error) { + return sendPipe(dst, data, W) } func SendConn(dst Sender, data []byte) (conn *UnixConn, err error) { diff --git a/pkg/beam/examples/beamsh/beamsh.go b/pkg/beam/examples/beamsh/beamsh.go index 3f258de332..808f038c68 100644 --- a/pkg/beam/examples/beamsh/beamsh.go +++ b/pkg/beam/examples/beamsh/beamsh.go @@ -257,12 +257,12 @@ func Handlers(sink beam.Sender) (*beam.UnixConn, error) { if handler == nil { return } - stdout, err := beam.SendPipe(conn, data.Empty().Set("cmd", "log", "stdout").Set("fromcmd", cmd...).Bytes()) + stdout, err := beam.SendRPipe(conn, data.Empty().Set("cmd", "log", "stdout").Set("fromcmd", cmd...).Bytes()) if err != nil { return } defer stdout.Close() - stderr, err := beam.SendPipe(conn, data.Empty().Set("cmd", "log", "stderr").Set("fromcmd", cmd...).Bytes()) + stderr, err := beam.SendRPipe(conn, data.Empty().Set("cmd", "log", "stderr").Set("fromcmd", cmd...).Bytes()) if err != nil { return } diff --git a/pkg/beam/examples/beamsh/builtins.go b/pkg/beam/examples/beamsh/builtins.go index cc94d2b5fb..3242237cc1 100644 --- a/pkg/beam/examples/beamsh/builtins.go +++ b/pkg/beam/examples/beamsh/builtins.go @@ -272,7 +272,7 @@ func CmdPrint(args []string, stdout, stderr io.Writer, in beam.Receiver, out bea } // Skip commands if a != nil && data.Message(payload).Get("cmd") == nil { - dup, err := beam.SendPipe(out, payload) + dup, err := beam.SendRPipe(out, payload) if err != nil { a.Close() return diff --git a/pkg/beam/router.go b/pkg/beam/router.go index fc41a8991b..15910e95b1 100644 --- a/pkg/beam/router.go +++ b/pkg/beam/router.go @@ -78,7 +78,7 @@ func (route *Route) Tee(dst Sender) *Route { return inner(payload, attachment) } // Setup the tee - w, err := SendPipe(dst, payload) + w, err := SendRPipe(dst, payload) if err != nil { return err }