Add binary stream functionality and helpers
Signed-off-by: Derek McGowan <derek@mcg.dev>
This commit is contained in:
Derek McGowan
210
pkg/transfer/streaming/stream.go
Normal file
210
pkg/transfer/streaming/stream.go
Normal file
@@ -0,0 +1,210 @@
|
||||
/*
|
||||
Copyright The containerd Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"crypto/rand"
|
||||
"encoding/base64"
|
||||
"errors"
|
||||
"fmt"
|
||||
"io"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
transferapi "github.com/containerd/containerd/api/types/transfer"
|
||||
"github.com/containerd/containerd/log"
|
||||
"github.com/containerd/containerd/pkg/streaming"
|
||||
"github.com/containerd/typeurl"
|
||||
)
|
||||
|
||||
const maxRead = 32 * 1024
|
||||
const windowSize = 2 * maxRead
|
||||
|
||||
var bufPool = &sync.Pool{
|
||||
New: func() interface{} {
|
||||
buffer := make([]byte, maxRead)
|
||||
return &buffer
|
||||
},
|
||||
}
|
||||
|
||||
func SendStream(ctx context.Context, r io.Reader, stream streaming.Stream) {
|
||||
window := make(chan int32)
|
||||
go func() {
|
||||
defer close(window)
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
any, err := stream.Recv()
|
||||
if err != nil {
|
||||
if !errors.Is(err, io.EOF) && !errors.Is(err, context.Canceled) {
|
||||
log.G(ctx).WithError(err).Error("send stream ended without EOF")
|
||||
}
|
||||
return
|
||||
}
|
||||
i, err := typeurl.UnmarshalAny(any)
|
||||
if err != nil {
|
||||
log.G(ctx).WithError(err).Error("failed to unmarshal stream object")
|
||||
continue
|
||||
}
|
||||
switch v := i.(type) {
|
||||
case *transferapi.WindowUpdate:
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case window <- v.Update:
|
||||
}
|
||||
default:
|
||||
log.G(ctx).Errorf("unexpected stream object of type %T", i)
|
||||
}
|
||||
}
|
||||
}()
|
||||
go func() {
|
||||
defer stream.Close()
|
||||
|
||||
buf := bufPool.Get().(*[]byte)
|
||||
defer bufPool.Put(buf)
|
||||
|
||||
var remaining int32
|
||||
|
||||
for {
|
||||
if remaining > 0 {
|
||||
// Don't wait for window update since there are remaining
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// TODO: Send error message on stream before close to allow remote side to return error
|
||||
return
|
||||
case update := <-window:
|
||||
remaining += update
|
||||
default:
|
||||
}
|
||||
} else {
|
||||
// Block until window updated
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
// TODO: Send error message on stream before close to allow remote side to return error
|
||||
return
|
||||
case update := <-window:
|
||||
remaining = update
|
||||
}
|
||||
}
|
||||
var max int32 = maxRead
|
||||
if max > remaining {
|
||||
max = remaining
|
||||
}
|
||||
b := (*buf)[:max]
|
||||
n, err := r.Read(b)
|
||||
if err != nil {
|
||||
if err != io.EOF {
|
||||
log.G(ctx).WithError(err).Errorf("failed to read stream source")
|
||||
// TODO: Send error message on stream before close to allow remote side to return error
|
||||
}
|
||||
return
|
||||
}
|
||||
remaining = remaining - int32(n)
|
||||
|
||||
data := &transferapi.Data{
|
||||
Data: b[:n],
|
||||
}
|
||||
any, err := typeurl.MarshalAny(data)
|
||||
if err != nil {
|
||||
log.G(ctx).WithError(err).Errorf("failed to marshal data for send")
|
||||
// TODO: Send error message on stream before close to allow remote side to return error
|
||||
return
|
||||
}
|
||||
if err := stream.Send(any); err != nil {
|
||||
log.G(ctx).WithError(err).Errorf("send failed")
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
}
|
||||
|
||||
func ReceiveStream(ctx context.Context, stream streaming.Stream) io.Reader {
|
||||
r, w := io.Pipe()
|
||||
go func() {
|
||||
defer stream.Close()
|
||||
var window int32
|
||||
for {
|
||||
var werr error
|
||||
if window < windowSize {
|
||||
update := &transferapi.WindowUpdate{
|
||||
Update: windowSize,
|
||||
}
|
||||
any, err := typeurl.MarshalAny(update)
|
||||
if err != nil {
|
||||
w.CloseWithError(fmt.Errorf("failed to marshal window update: %w", err))
|
||||
return
|
||||
}
|
||||
// check window update error after recv, stream may be complete
|
||||
if werr = stream.Send(any); werr == nil {
|
||||
window += windowSize
|
||||
} else if werr == io.EOF {
|
||||
// TODO: Why does send return EOF here
|
||||
werr = nil
|
||||
}
|
||||
}
|
||||
any, err := stream.Recv()
|
||||
if err != nil {
|
||||
if err == io.EOF {
|
||||
err = nil
|
||||
} else {
|
||||
err = fmt.Errorf("received failed: %w", err)
|
||||
}
|
||||
w.CloseWithError(err)
|
||||
return
|
||||
} else if werr != nil {
|
||||
// Try receive before erroring out
|
||||
w.CloseWithError(fmt.Errorf("failed to send window update: %w", werr))
|
||||
return
|
||||
}
|
||||
i, err := typeurl.UnmarshalAny(any)
|
||||
if err != nil {
|
||||
w.CloseWithError(fmt.Errorf("failed to unmarshal received object: %w", err))
|
||||
return
|
||||
}
|
||||
switch v := i.(type) {
|
||||
case *transferapi.Data:
|
||||
n, err := w.Write(v.Data)
|
||||
if err != nil {
|
||||
w.CloseWithError(fmt.Errorf("failed to unmarshal received object: %w", err))
|
||||
// Close will error out sender
|
||||
return
|
||||
}
|
||||
window = window - int32(n)
|
||||
// TODO: Handle error case
|
||||
default:
|
||||
log.G(ctx).Warnf("Ignoring unknown stream object of type %T", i)
|
||||
continue
|
||||
}
|
||||
}
|
||||
|
||||
}()
|
||||
|
||||
return r
|
||||
}
|
||||
|
||||
func GenerateID(prefix string) string {
|
||||
t := time.Now()
|
||||
var b [3]byte
|
||||
rand.Read(b[:])
|
||||
return fmt.Sprintf("%s-%d-%s", prefix, t.Nanosecond(), base64.URLEncoding.EncodeToString(b[:]))
|
||||
}
|
||||
178
pkg/transfer/streaming/stream_test.go
Normal file
178
pkg/transfer/streaming/stream_test.go
Normal file
@@ -0,0 +1,178 @@
|
||||
/*
|
||||
Copyright The containerd Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"io"
|
||||
"testing"
|
||||
|
||||
"github.com/containerd/containerd/pkg/streaming"
|
||||
"github.com/containerd/typeurl"
|
||||
)
|
||||
|
||||
func FuzzSendAndReceive(f *testing.F) {
|
||||
f.Add([]byte{})
|
||||
f.Add([]byte{0})
|
||||
f.Add(bytes.Repeat([]byte{0}, windowSize+1))
|
||||
f.Add([]byte("hello"))
|
||||
f.Add(bytes.Repeat([]byte("hello"), windowSize+1))
|
||||
f.Fuzz(func(t *testing.T, expected []byte) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
rs, ws := pipeStream()
|
||||
r, w := io.Pipe()
|
||||
SendStream(ctx, r, ws)
|
||||
or := ReceiveStream(ctx, rs)
|
||||
|
||||
go func() {
|
||||
io.Copy(w, bytes.NewBuffer(expected))
|
||||
w.Close()
|
||||
}()
|
||||
|
||||
actual, err := io.ReadAll(or)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(expected, actual) {
|
||||
t.Fatalf("received bytes are not equal\n\tactual: %v\n\texpected:%v", actual, expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
func FuzzSendAndReceiveChain(f *testing.F) {
|
||||
f.Add([]byte{})
|
||||
f.Add([]byte{0})
|
||||
f.Add(bytes.Repeat([]byte{0}, windowSize+1))
|
||||
f.Add([]byte("hello"))
|
||||
f.Add(bytes.Repeat([]byte("hello"), windowSize+1))
|
||||
f.Fuzz(func(t *testing.T, expected []byte) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
r, w := io.Pipe()
|
||||
|
||||
or := chainStreams(ctx, chainStreams(ctx, chainStreams(ctx, r)))
|
||||
|
||||
go func() {
|
||||
io.Copy(w, bytes.NewBuffer(expected))
|
||||
w.Close()
|
||||
}()
|
||||
|
||||
actual, err := io.ReadAll(or)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(expected, actual) {
|
||||
t.Fatalf("received bytes are not equal\n\tactual: %v\n\texpected:%v", actual, expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func FuzzWriter(f *testing.F) {
|
||||
f.Add([]byte{})
|
||||
f.Add([]byte{0})
|
||||
f.Add(bytes.Repeat([]byte{0}, windowSize+1))
|
||||
f.Add([]byte("hello"))
|
||||
f.Add(bytes.Repeat([]byte("hello"), windowSize+1))
|
||||
f.Fuzz(func(t *testing.T, expected []byte) {
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
defer cancel()
|
||||
|
||||
rs, ws := pipeStream()
|
||||
wc := WriteByteStream(ctx, ws)
|
||||
or := ReceiveStream(ctx, rs)
|
||||
|
||||
go func() {
|
||||
io.Copy(wc, bytes.NewBuffer(expected))
|
||||
wc.Close()
|
||||
}()
|
||||
|
||||
actual, err := io.ReadAll(or)
|
||||
if err != nil {
|
||||
t.Fatal(err)
|
||||
}
|
||||
|
||||
if !bytes.Equal(expected, actual) {
|
||||
t.Fatalf("received bytes are not equal\n\tactual: %v\n\texpected:%v", actual, expected)
|
||||
}
|
||||
})
|
||||
}
|
||||
|
||||
func chainStreams(ctx context.Context, r io.Reader) io.Reader {
|
||||
rs, ws := pipeStream()
|
||||
SendStream(ctx, r, ws)
|
||||
return ReceiveStream(ctx, rs)
|
||||
}
|
||||
|
||||
func pipeStream() (streaming.Stream, streaming.Stream) {
|
||||
r := make(chan typeurl.Any)
|
||||
rc := make(chan struct{})
|
||||
w := make(chan typeurl.Any)
|
||||
wc := make(chan struct{})
|
||||
|
||||
rs := &testStream{
|
||||
send: w,
|
||||
recv: r,
|
||||
closer: wc,
|
||||
remote: rc,
|
||||
}
|
||||
ws := &testStream{
|
||||
send: r,
|
||||
recv: w,
|
||||
closer: rc,
|
||||
remote: wc,
|
||||
}
|
||||
return rs, ws
|
||||
}
|
||||
|
||||
type testStream struct {
|
||||
send chan<- typeurl.Any
|
||||
recv <-chan typeurl.Any
|
||||
closer chan struct{}
|
||||
remote <-chan struct{}
|
||||
}
|
||||
|
||||
func (ts *testStream) Send(a typeurl.Any) error {
|
||||
select {
|
||||
case <-ts.remote:
|
||||
return io.ErrClosedPipe
|
||||
case ts.send <- a:
|
||||
}
|
||||
return nil
|
||||
}
|
||||
func (ts *testStream) Recv() (typeurl.Any, error) {
|
||||
select {
|
||||
case <-ts.remote:
|
||||
return nil, io.EOF
|
||||
case a := <-ts.recv:
|
||||
return a, nil
|
||||
}
|
||||
}
|
||||
|
||||
func (ts *testStream) Close() error {
|
||||
select {
|
||||
case <-ts.closer:
|
||||
return nil
|
||||
default:
|
||||
}
|
||||
close(ts.closer)
|
||||
return nil
|
||||
}
|
||||
130
pkg/transfer/streaming/writer.go
Normal file
130
pkg/transfer/streaming/writer.go
Normal file
@@ -0,0 +1,130 @@
|
||||
/*
|
||||
Copyright The containerd Authors.
|
||||
|
||||
Licensed under the Apache License, Version 2.0 (the "License");
|
||||
you may not use this file except in compliance with the License.
|
||||
You may obtain a copy of the License at
|
||||
|
||||
http://www.apache.org/licenses/LICENSE-2.0
|
||||
|
||||
Unless required by applicable law or agreed to in writing, software
|
||||
distributed under the License is distributed on an "AS IS" BASIS,
|
||||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
See the License for the specific language governing permissions and
|
||||
limitations under the License.
|
||||
*/
|
||||
|
||||
package streaming
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"sync/atomic"
|
||||
|
||||
transferapi "github.com/containerd/containerd/api/types/transfer"
|
||||
"github.com/containerd/containerd/log"
|
||||
"github.com/containerd/containerd/pkg/streaming"
|
||||
"github.com/containerd/typeurl"
|
||||
)
|
||||
|
||||
func WriteByteStream(ctx context.Context, stream streaming.Stream) io.WriteCloser {
|
||||
wbs := &writeByteStream{
|
||||
ctx: ctx,
|
||||
stream: stream,
|
||||
updated: make(chan struct{}, 1),
|
||||
}
|
||||
go func() {
|
||||
for {
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
default:
|
||||
}
|
||||
|
||||
any, err := stream.Recv()
|
||||
if err != nil {
|
||||
if !errors.Is(err, io.EOF) && !errors.Is(err, context.Canceled) {
|
||||
log.G(ctx).WithError(err).Error("send byte stream ended without EOF")
|
||||
}
|
||||
return
|
||||
}
|
||||
i, err := typeurl.UnmarshalAny(any)
|
||||
if err != nil {
|
||||
log.G(ctx).WithError(err).Error("failed to unmarshal stream object")
|
||||
continue
|
||||
}
|
||||
switch v := i.(type) {
|
||||
case *transferapi.WindowUpdate:
|
||||
atomic.AddInt32(&wbs.remaining, v.Update)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
return
|
||||
case wbs.updated <- struct{}{}:
|
||||
default:
|
||||
// Don't block if no writes are waiting
|
||||
}
|
||||
default:
|
||||
log.G(ctx).Errorf("unexpected stream object of type %T", i)
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
return wbs
|
||||
}
|
||||
|
||||
type writeByteStream struct {
|
||||
ctx context.Context
|
||||
stream streaming.Stream
|
||||
remaining int32
|
||||
updated chan struct{}
|
||||
}
|
||||
|
||||
func (wbs *writeByteStream) Write(p []byte) (n int, err error) {
|
||||
for len(p) > 0 {
|
||||
remaining := atomic.LoadInt32(&wbs.remaining)
|
||||
if remaining == 0 {
|
||||
// Don't wait for window update since there are remaining
|
||||
select {
|
||||
case <-wbs.ctx.Done():
|
||||
// TODO: Send error message on stream before close to allow remote side to return error
|
||||
err = io.ErrShortWrite
|
||||
return
|
||||
case <-wbs.updated:
|
||||
continue
|
||||
}
|
||||
}
|
||||
var max int32 = maxRead
|
||||
if max > int32(len(p)) {
|
||||
max = int32(len(p))
|
||||
}
|
||||
if max > remaining {
|
||||
max = remaining
|
||||
}
|
||||
// TODO: continue
|
||||
//remaining = remaining - int32(n)
|
||||
|
||||
data := &transferapi.Data{
|
||||
Data: p[:max],
|
||||
}
|
||||
var any typeurl.Any
|
||||
any, err = typeurl.MarshalAny(data)
|
||||
if err != nil {
|
||||
log.G(wbs.ctx).WithError(err).Errorf("failed to marshal data for send")
|
||||
// TODO: Send error message on stream before close to allow remote side to return error
|
||||
return
|
||||
}
|
||||
if err = wbs.stream.Send(any); err != nil {
|
||||
log.G(wbs.ctx).WithError(err).Errorf("send failed")
|
||||
return
|
||||
}
|
||||
n += int(max)
|
||||
p = p[max:]
|
||||
atomic.AddInt32(&wbs.remaining, -1*max)
|
||||
}
|
||||
return
|
||||
}
|
||||
|
||||
func (wbs *writeByteStream) Close() error {
|
||||
return wbs.stream.Close()
|
||||
}
|
||||
Reference in New Issue
Block a user