Merge pull request AliyunContainerService#854 from YaoZengzeng/exec-a…

…ttach

feature: implement exec and attach method of stream server

cri/stream/remotecommand/attach.go

@@ -1,15 +1,14 @@
 package remotecommand
 
 import (
-	"fmt"
 	"net/http"
 	"time"
 )
 
 // Attacher knows how to attach a running container in a pod.
 type Attacher interface {
 	// Attach attaches to the running container in the pod.
-	Attach() error
+	Attach(containerID string, streamOpts *Options, streams *Streams) error
 }
 
 // ServeAttach handles requests to attach to a container. After creating/receiving the required
@@ -22,9 +21,10 @@ func ServeAttach(w http.ResponseWriter, req *http.Request, attacher Attacher, co
 	}
 	defer ctx.conn.Close()
 
-	// Hardcode to pass CI, implement it later.
-	fmt.Fprintf(ctx.stdoutStream, "hello\n")
-
-	// Actuall it's a bug of cri-tools v1.0.0-alpha.0, workaround it.
-	time.Sleep(1 * time.Second)
+	attacher.Attach(container, streamOpts, &Streams{
+		StreamCh:     make(chan struct{}, 1),
+		StdinStream:  ctx.stdinStream,
+		StdoutStream: ctx.stdoutStream,
+		StderrStream: ctx.stderrStream,
+	})
 }

cri/stream/remotecommand/exec.go

@@ -1,15 +1,14 @@
 package remotecommand
 
 import (
-	"fmt"
 	"net/http"
 	"time"
 )
 
 // Executor knows how to execute a command in a container of the pod.
 type Executor interface {
 	// Exec executes a command in a container of the pod.
-	Exec() error
+	Exec(containerID string, cmd []string, streamOpts *Options, streams *Streams) error
 }
 
 // ServeExec handles requests to execute a command in a container. After
@@ -23,6 +22,9 @@ func ServeExec(w http.ResponseWriter, req *http.Request, executor Executor, cont
 	}
 	defer ctx.conn.Close()
 
-	// Hardcode to pass CI, implement it later.
-	fmt.Fprintf(ctx.stdoutStream, "hello\n")
+	executor.Exec(container, cmd, streamOpts, &Streams{
+		StdinStream:  ctx.stdinStream,
+		StdoutStream: ctx.stdoutStream,
+		StderrStream: ctx.stderrStream,
+	})
 }

cri/stream/remotecommand/httpstream.go

@@ -22,6 +22,16 @@ type Options struct {
 	TTY    bool
 }
 
+// Streams contains all the streams used to stdio for
+// remote command execution.
+type Streams struct {
+	// Notified from StreamCh if streams broken.
+	StreamCh     chan struct{}
+	StdinStream  io.ReadCloser
+	StdoutStream io.WriteCloser
+	StderrStream io.WriteCloser
+}
+
 // context contains the connection and streams used when
 // forwarding an attach or execute session into a container.
 type context struct {
@@ -82,6 +92,8 @@ func createHTTPStreamStreams(w http.ResponseWriter, req *http.Request, opts *Opt
 	case "":
 		logrus.Infof("Client did not request protocol negotiation. Falling back to %q", constant.StreamProtocolV1Name)
 		fallthrough
+	case constant.StreamProtocolV2Name:
+		handler = &v2ProtocolHandler{}
 	case constant.StreamProtocolV1Name:
 		handler = &v1ProtocolHandler{}
 	}
@@ -129,6 +141,50 @@ type protocolHandler interface {
 	waitForStreams(streams <-chan streamAndReply, expectedStreams int, expired <-chan time.Time) (*context, error)
 }
 
+// v2ProtocolHandler implements the V2 protocol version for streaming command execution.
+type v2ProtocolHandler struct{}
+
+func (*v2ProtocolHandler) waitForStreams(streams <-chan streamAndReply, expectedStreams int, expired <-chan time.Time) (*context, error) {
+	ctx := &context{}
+	receivedStreams := 0
+	replyChan := make(chan struct{})
+	stop := make(chan struct{})
+	defer close(stop)
+WaitForStreams:
+	for {
+		select {
+		case stream := <-streams:
+			streamType := stream.Headers().Get(constant.StreamType)
+			switch streamType {
+			case constant.StreamTypeError:
+				go waitStreamReply(stream.replySent, replyChan, stop)
+			case constant.StreamTypeStdin:
+				ctx.stdinStream = stream
+				go waitStreamReply(stream.replySent, replyChan, stop)
+			case constant.StreamTypeStdout:
+				ctx.stdoutStream = stream
+				go waitStreamReply(stream.replySent, replyChan, stop)
+			case constant.StreamTypeStderr:
+				ctx.stderrStream = stream
+				go waitStreamReply(stream.replySent, replyChan, stop)
+			default:
+				logrus.Errorf("Unexpected stream type: %q", streamType)
+			}
+		case <-replyChan:
+			receivedStreams++
+			if receivedStreams == expectedStreams {
+				break WaitForStreams
+			}
+		case <-expired:
+			// TODO find a way to return the error to the user. Maybe use a separate
+			// stream to report errors?
+			return nil, fmt.Errorf("timed out waiting for client to create streams")
+		}
+	}
+
+	return ctx, nil
+}
+
 // v1ProtocolHandler implements the V1 protocol version for streaming command execution.
 type v1ProtocolHandler struct{}
 

cri/stream/server.go

@@ -30,7 +30,7 @@ const (
 // TODO: StreamProtocolV2Name, StreamProtocolV3Name, StreamProtocolV4Name support.
 
 // SupportedStreamingProtocols is the streaming protocols which server supports.
-var SupportedStreamingProtocols = []string{constant.StreamProtocolV1Name}
+var SupportedStreamingProtocols = []string{constant.StreamProtocolV1Name, constant.StreamProtocolV2Name}
 
 // SupportedPortForwardProtocols is the portforward protocols which server supports.
 var SupportedPortForwardProtocols = []string{constant.PortForwardProtocolV1Name}
@@ -53,10 +53,10 @@ type Server interface {
 // Runtime is the interface to execute the commands and provide the streams.
 type Runtime interface {
 	// Exec executes the command in pod.
-	Exec() error
+	Exec(containerID string, cmd []string, streamOpts *remotecommand.Options, streams *remotecommand.Streams) error
 
 	// Attach attaches to pod.
-	Attach() error
+	Attach(containerID string, streamOpts *remotecommand.Options, streams *remotecommand.Streams) error
 
 	// PortForward forward port to pod.
 	PortForward(name string, port int32, stream io.ReadWriteCloser) error

daemon/containerio/container_io.go

@@ -50,6 +50,19 @@ func NewIO(opt *Option) *IO {
 	}
 }
 
+// AddBackend adds more backends to container's stdio.
+func (io *IO) AddBackend(opt *Option) {
+	backends := createBackend(opt)
+
+	for t, s := range map[stdioType]*ContainerIO{
+		stdin:  io.Stdin,
+		stdout: io.Stdout,
+		stderr: io.Stderr,
+	} {
+		s.add(opt, t, backends)
+	}
+}
+
 // Close closes the container's io.
 func (io *IO) Close() error {
 	io.Stderr.Close()
@@ -61,30 +74,59 @@ func (io *IO) Close() error {
 // ContainerIO used to control the container's stdio.
 type ContainerIO struct {
 	Option
-	backends map[string]containerBackend
+	backends []containerBackend
 	total    int64
 	typ      stdioType
 	closed   bool
+	// The stdin of all backends should put into ring first.
+	ring *ringbuff.RingBuff
+}
+
+func (cio *ContainerIO) add(opt *Option, typ stdioType, backends map[string]containerBackend) {
+	if typ == stdin {
+		for _, b := range backends {
+			if b.backend.Name() == opt.stdinBackend {
+				cio.backends = append(cio.backends, b)
+				go func(b containerBackend) {
+					cio.converge(b.backend.Name(), opt.id, b.backend.In())
+					b.backend.Close()
+				}(b)
+				break
+			}
+		}
+	} else {
+		for _, b := range backends {
+			cio.backends = append(cio.backends, b)
+		}
+	}
 }
 
 func create(opt *Option, typ stdioType, backends map[string]containerBackend) *ContainerIO {
 	io := &ContainerIO{
-		backends: backends,
-		total:    0,
-		typ:      typ,
-		closed:   false,
-		Option:   *opt,
+		total:  0,
+		typ:    typ,
+		closed: false,
+		Option: *opt,
 	}
 
 	if typ == stdin {
-		io.backends = make(map[string]containerBackend)
-
+		io.ring = ringbuff.New(10)
 		for _, b := range backends {
 			if b.backend.Name() == opt.stdinBackend {
-				io.backends[opt.stdinBackend] = b
+				io.backends = append(io.backends, b)
+				go func(b containerBackend) {
+					// For backend with stdin, close it if stdin finished.
+					io.converge(b.backend.Name(), opt.id, b.backend.In())
+					b.backend.Close()
+					b.ring.Close()
+				}(b)
 				break
 			}
 		}
+	} else {
+		for _, b := range backends {
+			io.backends = append(io.backends, b)
+		}
 	}
 
 	return io
@@ -124,69 +166,70 @@ func createBackend(opt *Option) map[string]containerBackend {
 }
 
 // OpenStdin returns open container's stdin or not.
-func (io *ContainerIO) OpenStdin() bool {
-	if io.typ != stdin {
+func (cio *ContainerIO) OpenStdin() bool {
+	if cio.typ != stdin {
 		return false
 	}
-	if io.closed {
+	if cio.closed {
 		return false
 	}
-	return len(io.backends) != 0
+	return len(cio.backends) != 0
 }
 
 // Read implements the standard Read interface.
-func (io *ContainerIO) Read(p []byte) (int, error) {
-	if io.typ != stdin {
-		return 0, fmt.Errorf("invalid container io type: %s, id: %s", io.typ, io.id)
+func (cio *ContainerIO) Read(p []byte) (int, error) {
+	if cio.typ != stdin {
+		return 0, fmt.Errorf("invalid container io type: %s, id: %s", cio.typ, cio.id)
 	}
-	if io.closed {
+	if cio.closed {
 		return 0, fmt.Errorf("container io is closed")
 	}
 
-	if len(io.backends) == 0 {
-		block := make(chan struct{})
-		<-block
+	value, _ := cio.ring.Pop()
+	data, ok := value.([]byte)
+	if !ok {
+		return 0, nil
 	}
+	n := copy(p, data)
 
-	backend := io.backends[io.stdinBackend]
-
-	return backend.backend.In().Read(p)
+	return n, nil
 }
 
 // Write implements the standard Write interface.
-func (io *ContainerIO) Write(data []byte) (int, error) {
-	if io.typ == stdin {
-		return 0, fmt.Errorf("invalid container io type: %s, id: %s", io.typ, io.id)
+func (cio *ContainerIO) Write(data []byte) (int, error) {
+	if cio.typ == stdin {
+		return 0, fmt.Errorf("invalid container io type: %s, id: %s", cio.typ, cio.id)
 	}
-	if io.closed {
+	if cio.closed {
 		return 0, fmt.Errorf("container io is closed")
 	}
 
-	if io.typ == discard {
+	if cio.typ == discard {
 		return len(data), nil
 	}
 
-	for _, b := range io.backends {
+	for _, b := range cio.backends {
 		if cover := b.ring.Push(data); cover {
-			logrus.Warnf("cover data, backend: %s, id: %s", b.backend.Name(), io.id)
+			logrus.Warnf("cover data, backend: %s, id: %s", b.backend.Name(), cio.id)
 		}
 	}
 
 	return len(data), nil
 }
 
 // Close implements the standard Close interface.
-func (io *ContainerIO) Close() error {
-	for name, b := range io.backends {
+func (cio *ContainerIO) Close() error {
+	for _, b := range cio.backends {
 		// we need to close ringbuf before close backend, because close ring will flush
 		// the remain data into backend.
+		name := b.backend.Name()
 		b.ring.Close()
 		b.backend.Close()
 
-		logrus.Infof("close containerio backend: %s, id: %s", name, io.id)
+		logrus.Infof("close containerio backend: %s, id: %s", name, cio.id)
 	}
 
-	io.closed = true
+	cio.closed = true
 	return nil
 }
 
@@ -215,3 +258,24 @@ func subscribe(name, id string, ring *ringbuff.RingBuff, out io.Writer) {
 
 	logrus.Infof("finished to subscribe io, backend: %s, id: %s", name, id)
 }
+
+// converge be called in a goroutine.
+func (cio *ContainerIO) converge(name, id string, in io.Reader) {
+	// TODO: we should implement this function more elegant and robust.
+	logrus.Infof("start to converge io, backend: %s, id: %s", name, id)
+
+	data := make([]byte, 128)
+	for {
+		n, err := in.Read(data)
+		if err != nil {
+			logrus.Errorf("failed to read from backend: %s, id: %s, %v", name, id, err)
+			break
+		}
+		cover := cio.ring.Push(data[:n])
+		if cover {
+			logrus.Warnf("cover data, backend: %s, id: %s", name, id)
+		}
+	}
+
+	logrus.Infof("finished to converge io, backend: %s, id: %s", name, id)
+}