Implicitly discard the input to drain the reader

Signed-off-by: Derek McGowan <derek@mcg.dev>
This commit is contained in:
Kazuyoshi Kato 2022-06-02 03:53:11 +00:00 committed by Derek McGowan
parent 49ca87d727
commit 40aa4f3f1b
No known key found for this signature in database
GPG Key ID: F58C5D0A4405ACDB
2 changed files with 23 additions and 14 deletions

View File

@ -19,7 +19,6 @@ package server
import ( import (
"bytes" "bytes"
"context" "context"
"errors"
"fmt" "fmt"
"io" "io"
"syscall" "syscall"
@ -44,11 +43,9 @@ type cappedWriter struct {
remain int remain int
} }
var errNoRemain = errors.New("no more space to write")
func (cw *cappedWriter) Write(p []byte) (int, error) { func (cw *cappedWriter) Write(p []byte) (int, error) {
if cw.remain <= 0 { if cw.remain <= 0 {
return 0, errNoRemain return len(p), nil
} }
end := cw.remain end := cw.remain
@ -61,26 +58,35 @@ func (cw *cappedWriter) Write(p []byte) (int, error) {
if err != nil { if err != nil {
return written, err return written, err
} }
if written < len(p) { return len(p), nil
return written, errNoRemain
}
return written, nil
} }
func (cw *cappedWriter) Close() error { func (cw *cappedWriter) Close() error {
return cw.w.Close() return cw.w.Close()
} }
func (cw *cappedWriter) isFull() bool {
return cw.remain <= 0
}
// ExecSync executes a command in the container, and returns the stdout output. // ExecSync executes a command in the container, and returns the stdout output.
// If command exits with a non-zero exit code, an error is returned. // If command exits with a non-zero exit code, an error is returned.
func (c *criService) ExecSync(ctx context.Context, r *runtime.ExecSyncRequest) (*runtime.ExecSyncResponse, error) { func (c *criService) ExecSync(ctx context.Context, r *runtime.ExecSyncRequest) (*runtime.ExecSyncResponse, error) {
const maxStreamSize = 1024 * 1024 * 16 const maxStreamSize = 1024 * 1024 * 16
var stdout, stderr bytes.Buffer var stdout, stderr bytes.Buffer
// cappedWriter truncates the output. In that case, the size of
// the ExecSyncResponse will hit the CRI plugin's gRPC response limit.
// Thus the callers outside of the containerd process (e.g. Kubelet) never see
// the truncated output.
cout := &cappedWriter{w: cioutil.NewNopWriteCloser(&stdout), remain: maxStreamSize}
cerr := &cappedWriter{w: cioutil.NewNopWriteCloser(&stderr), remain: maxStreamSize}
exitCode, err := c.execInContainer(ctx, r.GetContainerId(), execOptions{ exitCode, err := c.execInContainer(ctx, r.GetContainerId(), execOptions{
cmd: r.GetCmd(), cmd: r.GetCmd(),
stdout: &cappedWriter{w: cioutil.NewNopWriteCloser(&stdout), remain: maxStreamSize}, stdout: cout,
stderr: &cappedWriter{w: cioutil.NewNopWriteCloser(&stderr), remain: maxStreamSize}, stderr: cerr,
timeout: time.Duration(r.GetTimeout()) * time.Second, timeout: time.Duration(r.GetTimeout()) * time.Second,
}) })
if err != nil { if err != nil {

View File

@ -33,12 +33,15 @@ func TestCWWrite(t *testing.T) {
assert.Equal(t, 5, n) assert.Equal(t, 5, n)
n, err = cw.Write([]byte("helloworld")) n, err = cw.Write([]byte("helloworld"))
assert.Equal(t, []byte("hellohello"), buf.Bytes(), "partial write") assert.NoError(t, err, "no errors even it hits the cap")
assert.Equal(t, 5, n) assert.Equal(t, 10, n, "no indication of partial write")
assert.ErrorIs(t, err, errNoRemain) assert.True(t, cw.isFull())
assert.Equal(t, []byte("hellohello"), buf.Bytes(), "the underlying writer is capped")
_, err = cw.Write([]byte("world")) _, err = cw.Write([]byte("world"))
assert.ErrorIs(t, err, errNoRemain) assert.NoError(t, err)
assert.True(t, cw.isFull())
assert.Equal(t, []byte("hellohello"), buf.Bytes(), "the underlying writer is capped")
} }
func TestCWClose(t *testing.T) { func TestCWClose(t *testing.T) {