demultiplex exec output into stdout / stderr

cluster/calcium/control_test.go

@@ -52,13 +52,13 @@ func TestControlStart(t *testing.T) {
 		assert.Error(t, r.Error)
 	}
 	engine.On("VirtualizationStart", mock.Anything, mock.Anything).Return(nil)
-	// failed by ExecCreate
+	// failed by Execute
 	hook := &types.Hook{
 		AfterStart: []string{"cmd1", "cmd2"},
 	}
 	workload.Hook = hook
 	workload.Hook.Force = false
-	engine.On("ExecCreate", mock.Anything, mock.Anything, mock.Anything).Return("", types.ErrNilEngine).Times(3)
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("", nil, nil, nil, types.ErrNilEngine).Times(3)
 	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadStart, false)
 	assert.NoError(t, err)
 	for r := range ch {
@@ -72,16 +72,8 @@ func TestControlStart(t *testing.T) {
 		assert.Error(t, r.Error)
 		assert.Equal(t, r.WorkloadID, "id1")
 	}
-	engine.On("ExecCreate", mock.Anything, mock.Anything, mock.Anything).Return("eid", nil)
-	// failed by ExecAttach
-	engine.On("ExecAttach", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(nil, nil, types.ErrNilEngine).Once()
-	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadStart, false)
-	assert.NoError(t, err)
-	for r := range ch {
-		assert.Error(t, r.Error)
-	}
 	data := ioutil.NopCloser(bytes.NewBufferString("output"))
-	engine.On("ExecAttach", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(data, nil, nil).Twice()
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("eid", data, nil, nil, nil).Times(4)
 	// failed by ExecExitCode
 	engine.On("ExecExitCode", mock.Anything, mock.Anything).Return(-1, types.ErrNilEngine).Once()
 	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadStart, false)
@@ -98,7 +90,6 @@ func TestControlStart(t *testing.T) {
 	}
 	// exitCode is 0
 	engine.On("ExecExitCode", mock.Anything, mock.Anything).Return(0, nil)
-	engine.On("ExecAttach", mock.Anything, mock.Anything, mock.Anything, mock.Anything).Return(ioutil.NopCloser(bytes.NewBufferString("succ")), nil, nil)
 	ch, err = c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadStart, false)
 	assert.NoError(t, err)
 	for r := range ch {
@@ -128,7 +119,7 @@ func TestControlStop(t *testing.T) {
 	}
 	workload.Hook = hook
 	workload.Hook.Force = true
-	engine.On("ExecCreate", mock.Anything, mock.Anything, mock.Anything).Return("", types.ErrNilEngine)
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("", nil, nil, nil, types.ErrNilEngine)
 	ch, err := c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadStop, false)
 	assert.NoError(t, err)
 	for r := range ch {
@@ -173,7 +164,7 @@ func TestControlRestart(t *testing.T) {
 	}
 	workload.Hook = hook
 	workload.Hook.Force = true
-	engine.On("ExecCreate", mock.Anything, mock.Anything, mock.Anything).Return("", types.ErrNilEngine)
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return("", nil, nil, nil, types.ErrNilEngine)
 	ch, err := c.ControlWorkload(ctx, []string{"id1"}, cluster.WorkloadRestart, false)
 	assert.NoError(t, err)
 	for r := range ch {

cluster/calcium/execute.go

@@ -1,6 +1,7 @@
 package calcium
 
 import (
+	"bufio"
 	"context"
 	"strconv"
 
@@ -14,19 +15,18 @@ func (c *Calcium) ExecuteWorkload(ctx context.Context, opts *types.ExecuteWorklo
 	ch := make(chan *types.AttachWorkloadMessage)
 
 	go func() {
-		defer close(ch)
 		var err error
-		responses := []string{}
+
 		defer func() {
-			for _, resp := range responses {
-				msg := &types.AttachWorkloadMessage{WorkloadID: opts.WorkloadID, Data: []byte(resp)}
-				ch <- msg
+			if err != nil {
+				ch <- &types.AttachWorkloadMessage{WorkloadID: opts.WorkloadID, Data: []byte(err.Error())}
 			}
+			close(ch)
 		}()
 
 		workload, err := c.GetWorkload(ctx, opts.WorkloadID)
 		if err != nil {
-			responses = append(responses, err.Error())
+			log.Errorf("[ExecuteWorkload] Failed to get wordload: %v", err)
 			return
 		}
 
@@ -41,20 +41,22 @@ func (c *Calcium) ExecuteWorkload(ctx context.Context, opts *types.ExecuteWorklo
 			Detach:       false,
 		}
 
-		execID, outStream, inStream, err := workload.Engine.Execute(ctx, opts.WorkloadID, execConfig)
+		execID, stdout, stderr, inStream, err := workload.Engine.Execute(ctx, opts.WorkloadID, execConfig)
 		if err != nil {
 			log.Errorf("[ExecuteWorkload] Failed to attach execID: %v", err)
 			return
 		}
 
+		splitFunc, split := bufio.ScanLines, byte('\n')
 		if opts.OpenStdin {
 			processVirtualizationInStream(ctx, inStream, inCh, func(height, width uint) error {
 				return workload.Engine.ExecResize(ctx, execID, height, width)
 			})
+			splitFunc, split = bufio.ScanBytes, byte(0)
 		}
 
-		for data := range processVirtualizationOutStream(ctx, outStream) {
-			ch <- &types.AttachWorkloadMessage{WorkloadID: opts.WorkloadID, Data: data}
+		for m := range processStdStream(ctx, stdout, stderr, splitFunc, split) {
+			ch <- &types.AttachWorkloadMessage{WorkloadID: opts.WorkloadID, Data: m.Data, StdStreamType: m.StdStreamType}
 		}
 
 		execCode, err := workload.Engine.ExecExitCode(ctx, execID)

cluster/calcium/execute_test.go

@@ -41,13 +41,13 @@ func TestExecuteWorkload(t *testing.T) {
 	}
 	store.On("GetWorkload", mock.Anything, mock.Anything).Return(workload, nil)
 	// failed by Execute
-	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(ID, nil, nil, types.ErrCannotGetEngine).Once()
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(ID, nil, nil, nil, types.ErrCannotGetEngine).Once()
 	ch = c.ExecuteWorkload(ctx, &types.ExecuteWorkloadOptions{WorkloadID: ID}, nil)
 	for ac := range ch {
 		assert.Equal(t, ac.WorkloadID, ID)
 	}
 	buf := ioutil.NopCloser(bytes.NewBufferString(`echo 1\n`))
-	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(ID, buf, nil, nil).Twice()
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(ID, buf, nil, nil, nil).Twice()
 	// failed by ExecExitCode
 	engine.On("ExecExitCode", mock.Anything, mock.Anything).Return(-1, types.ErrCannotGetEngine).Once()
 	ch = c.ExecuteWorkload(ctx, &types.ExecuteWorkloadOptions{WorkloadID: ID}, nil)
@@ -67,7 +67,7 @@ func TestExecuteWorkload(t *testing.T) {
 	assert.Contains(t, string(data), "0")
 	inChan := make(chan []byte)
 	inS := &inStream{bytes.NewBufferString("")}
-	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(ID, buf, inS, nil)
+	engine.On("Execute", mock.Anything, mock.Anything, mock.Anything).Return(ID, buf, nil, inS, nil)
 	ch = c.ExecuteWorkload(ctx, &types.ExecuteWorkloadOptions{WorkloadID: ID, OpenStdin: true}, inChan)
 	inChan <- []byte("a")
 	inChan <- escapeCommand

cluster/calcium/helper.go

@@ -6,7 +6,9 @@ import (
 	"fmt"
 	"io"
 	"io/ioutil"
+	"sync"
 
+	"github.com/pkg/errors"
 	"github.com/projecteru2/core/engine"
 	enginetypes "github.com/projecteru2/core/engine/types"
 	"github.com/projecteru2/core/log"
@@ -34,18 +36,13 @@ func execuateInside(ctx context.Context, client engine.API, ID, cmd, user string
 		AttachStdout: true,
 	}
 	b := []byte{}
-	execID, err := client.ExecCreate(ctx, ID, execConfig)
+	execID, stdout, stderr, _, err := client.Execute(ctx, ID, execConfig)
 	if err != nil {
-		return b, err
+		return nil, errors.WithStack(err)
 	}
 
-	outStream, _, err := client.ExecAttach(ctx, execID, false)
-	if err != nil {
-		return b, err
-	}
-
-	for data := range processVirtualizationOutStream(ctx, outStream) {
-		b = append(b, data...)
+	for m := range processStdStream(ctx, stdout, stderr, bufio.ScanLines, byte('\n')) {
+		b = append(b, m.Data...)
 	}
 
 	exitCode, err := client.ExecExitCode(ctx, execID)
@@ -176,16 +173,25 @@ func rawProcessVirtualizationInStream(
 func processVirtualizationOutStream(
 	_ context.Context,
 	outStream io.ReadCloser,
+	splitFunc bufio.SplitFunc,
+	split byte,
+
 ) <-chan []byte {
 	outCh := make(chan []byte)
 	go func() {
-		defer outStream.Close()
 		defer close(outCh)
+		if outStream == nil {
+			return
+		}
+		defer outStream.Close()
 		scanner := bufio.NewScanner(outStream)
-		scanner.Split(bufio.ScanBytes)
+		scanner.Split(splitFunc)
 		for scanner.Scan() {
-			b := scanner.Bytes()
-			outCh <- b
+			bs := scanner.Bytes()
+			if split != 0 {
+				bs = append(bs, split)
+			}
+			outCh <- bs
 		}
 		if err := scanner.Err(); err != nil {
 			log.Errorf("[processVirtualizationOutStream] failed to read output from output stream: %v", err)
@@ -217,3 +223,32 @@ func processBuildImageStream(reader io.ReadCloser) chan *types.BuildImageMessage
 	}()
 	return ch
 }
+
+func processStdStream(ctx context.Context, stdout, stderr io.ReadCloser, splitFunc bufio.SplitFunc, split byte) chan types.StdStreamMessage {
+	ch := make(chan types.StdStreamMessage)
+
+	wg := sync.WaitGroup{}
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for data := range processVirtualizationOutStream(ctx, stdout, splitFunc, split) {
+			ch <- types.StdStreamMessage{Data: data, StdStreamType: types.Stdout}
+		}
+	}()
+
+	wg.Add(1)
+	go func() {
+		defer wg.Done()
+		for data := range processVirtualizationOutStream(ctx, stderr, splitFunc, split) {
+			ch <- types.StdStreamMessage{Data: data, StdStreamType: types.Stderr}
+		}
+	}()
+
+	go func() {
+		defer close(ch)
+		wg.Wait()
+	}()
+
+	return ch
+}

cluster/calcium/lambda.go

@@ -1,6 +1,7 @@
 package calcium
 
 import (
+	"bufio"
 	"context"
 	"io"
 	"strconv"
@@ -57,17 +58,19 @@ func (c *Calcium) RunAndWait(ctx context.Context, opts *types.DeployOptions, inC
 				return
 			}
 
-			var outStream io.ReadCloser
-			if outStream, err = workload.Engine.VirtualizationLogs(ctx, &enginetypes.VirtualizationLogStreamOptions{
+			var stdout, stderr io.ReadCloser
+			if stdout, stderr, err = workload.Engine.VirtualizationLogs(ctx, &enginetypes.VirtualizationLogStreamOptions{
 				ID: message.WorkloadID, Follow: true, Stdout: true, Stderr: true}); err != nil {
 				log.Errorf("[RunAndWait] Can't fetch log of workload %s error %v", message.WorkloadID, err)
 				return
 			}
 
+			splitFunc, split := bufio.ScanLines, byte('\n')
+
 			// use attach if use stdin
 			if opts.OpenStdin {
 				var inStream io.WriteCloser
-				outStream, inStream, err = workload.Engine.VirtualizationAttach(ctx, message.WorkloadID, true, true)
+				stdout, stderr, inStream, err = workload.Engine.VirtualizationAttach(ctx, message.WorkloadID, true, true)
 				if err != nil {
 					log.Errorf("[RunAndWait] Can't attach workload %s error %v", message.WorkloadID, err)
 					return
@@ -76,10 +79,16 @@ func (c *Calcium) RunAndWait(ctx context.Context, opts *types.DeployOptions, inC
 				processVirtualizationInStream(ctx, inStream, inCh, func(height, width uint) error {
 					return workload.Engine.VirtualizationResize(ctx, message.WorkloadID, height, width)
 				})
+
+				splitFunc, split = bufio.ScanBytes, byte(0)
 			}
 
-			for data := range processVirtualizationOutStream(ctx, outStream) {
-				runMsgCh <- &types.AttachWorkloadMessage{WorkloadID: message.WorkloadID, Data: data}
+			for m := range processStdStream(ctx, stdout, stderr, splitFunc, split) {
+				runMsgCh <- &types.AttachWorkloadMessage{
+					WorkloadID:    message.WorkloadID,
+					Data:          m.Data,
+					StdStreamType: m.StdStreamType,
+				}
 			}
 
 			// wait and forward exitcode

cluster/calcium/log.go

@@ -19,7 +19,7 @@ func (c *Calcium) LogStream(ctx context.Context, opts *types.LogStreamOptions) (
 			return
 		}
 
-		resp, err := workload.Engine.VirtualizationLogs(ctx, &enginetypes.VirtualizationLogStreamOptions{
+		stdout, stderr, err := workload.Engine.VirtualizationLogs(ctx, &enginetypes.VirtualizationLogStreamOptions{
 			ID:     opts.ID,
 			Tail:   opts.Tail,
 			Since:  opts.Since,
@@ -33,11 +33,10 @@ func (c *Calcium) LogStream(ctx context.Context, opts *types.LogStreamOptions) (
 			return
 		}
 
-		scanner := bufio.NewScanner(resp)
-		for scanner.Scan() {
-			data := scanner.Bytes()
-			ch <- &types.LogStreamMessage{ID: opts.ID, Data: data}
+		for m := range processStdStream(ctx, stdout, stderr, bufio.ScanLines, byte('\n')) {
+			ch <- &types.LogStreamMessage{ID: opts.ID, Data: m.Data, StdStreamType: m.StdStreamType}
 		}
 	}()
+
 	return ch, nil
 }

cluster/calcium/log_test.go

@@ -36,15 +36,15 @@ func TestLogStream(t *testing.T) {
 	}
 	store.On("GetWorkload", mock.Anything, mock.Anything).Return(workload, nil)
 	// failed by VirtualizationLogs
-	engine.On("VirtualizationLogs", mock.Anything, mock.Anything).Return(nil, types.ErrNodeExist).Once()
+	engine.On("VirtualizationLogs", mock.Anything, mock.Anything).Return(nil, nil, types.ErrNodeExist).Once()
 	ch, err = c.LogStream(ctx, opts)
 	assert.NoError(t, err)
 	for c := range ch {
 		assert.Equal(t, c.ID, ID)
 		assert.Empty(t, c.Data)
 	}
 	reader := bytes.NewBufferString("aaaa\nbbbb\n")
-	engine.On("VirtualizationLogs", mock.Anything, mock.Anything).Return(ioutil.NopCloser(reader), nil)
+	engine.On("VirtualizationLogs", mock.Anything, mock.Anything).Return(ioutil.NopCloser(reader), nil, nil)
 	// success
 	ch, err = c.LogStream(ctx, opts)
 	assert.NoError(t, err)

engine/docker/container.go

@@ -270,7 +270,7 @@ func (e *Engine) VirtualizationInspect(ctx context.Context, ID string) (*enginet
 }
 
 // VirtualizationLogs show virtualization logs
-func (e *Engine) VirtualizationLogs(ctx context.Context, opts *enginetypes.VirtualizationLogStreamOptions) (io.ReadCloser, error) {
+func (e *Engine) VirtualizationLogs(ctx context.Context, opts *enginetypes.VirtualizationLogStreamOptions) (stdout, stderr io.ReadCloser, err error) {
 	logsOpts := dockertypes.ContainerLogsOptions{
 		ShowStdout: opts.Stdout,
 		ShowStderr: opts.Stderr,
@@ -281,13 +281,17 @@ func (e *Engine) VirtualizationLogs(ctx context.Context, opts *enginetypes.Virtu
 	}
 	resp, err := e.client.ContainerLogs(ctx, opts.ID, logsOpts)
 	if err != nil {
-		return nil, err
+		return nil, nil, err
+	}
+	if !opts.Stderr {
+		return ioutil.NopCloser(mergeStream(resp)), nil, nil
 	}
-	return ioutil.NopCloser(mergeStream(resp)), nil
+	stdout, stderr = e.demultiplexStdStream(resp)
+	return stdout, stderr, nil
 }
 
 // VirtualizationAttach attach to a virtualization
-func (e *Engine) VirtualizationAttach(ctx context.Context, ID string, stream, stdin bool) (io.ReadCloser, io.WriteCloser, error) {
+func (e *Engine) VirtualizationAttach(ctx context.Context, ID string, stream, stdin bool) (stdout, stderr io.ReadCloser, _ io.WriteCloser, err error) {
 	opts := dockertypes.ContainerAttachOptions{
 		Stream: stream,
 		Stdin:  stdin,
@@ -297,9 +301,13 @@ func (e *Engine) VirtualizationAttach(ctx context.Context, ID string, stream, st
 	}
 	resp, err := e.client.ContainerAttach(ctx, ID, opts)
 	if err != nil {
-		return nil, nil, err
+		return nil, nil, nil, err
+	}
+	if stdin {
+		return ioutil.NopCloser(resp.Reader), nil, resp.Conn, nil
 	}
-	return ioutil.NopCloser(resp.Reader), resp.Conn, nil
+	stdout, stderr = e.demultiplexStdStream(resp.Reader)
+	return stdout, stderr, resp.Conn, nil
 }
 
 // VirtualizationResize resizes remote terminal