move implementation to internal/ directory (#828)
This commit is contained in:
Christian Schwarz
committed by
GitHub
GitHub
parent
b9b9ad10cf
commit
908807bd59
@@ -0,0 +1,289 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"strings"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"unicode/utf8"
|
||||
|
||||
"github.com/zrepl/zrepl/internal/logger"
|
||||
"github.com/zrepl/zrepl/internal/rpc/dataconn/base2bufpool"
|
||||
"github.com/zrepl/zrepl/internal/rpc/dataconn/frameconn"
|
||||
"github.com/zrepl/zrepl/internal/rpc/dataconn/heartbeatconn"
|
||||
)
|
||||
|
||||
type Logger = logger.Logger
|
||||
|
||||
type contextKey int
|
||||
|
||||
const (
|
||||
contextKeyLogger contextKey = 1 + iota
|
||||
)
|
||||
|
||||
func WithLogger(ctx context.Context, log Logger) context.Context {
|
||||
return context.WithValue(ctx, contextKeyLogger, log)
|
||||
}
|
||||
|
||||
//nolint:deadcode,unused
|
||||
func getLog(ctx context.Context) Logger {
|
||||
log, ok := ctx.Value(contextKeyLogger).(Logger)
|
||||
if !ok {
|
||||
log = logger.NewNullLogger()
|
||||
}
|
||||
return log
|
||||
}
|
||||
|
||||
// Frame types used by this package.
|
||||
// 4 MSBs are reserved for frameconn, next 4 MSB for heartbeatconn, next 4 MSB for us.
|
||||
const (
|
||||
StreamErrTrailer uint32 = 1 << (16 + iota)
|
||||
End
|
||||
// max 16
|
||||
)
|
||||
|
||||
// NOTE: make sure to add a tests for each frame type that checks
|
||||
// whether it is heartbeatconn.IsPublicFrameType()
|
||||
|
||||
// Check whether the given frame type is allowed to be used by
|
||||
// consumers of this package. Intended for use in unit tests.
|
||||
func IsPublicFrameType(ft uint32) bool {
|
||||
return frameconn.IsPublicFrameType(ft) && heartbeatconn.IsPublicFrameType(ft) && ((0xf<<16)&ft == 0)
|
||||
}
|
||||
|
||||
const FramePayloadShift = 19
|
||||
|
||||
var bufpool = base2bufpool.New(FramePayloadShift, FramePayloadShift, base2bufpool.Panic)
|
||||
|
||||
// if sendStream returns an error, that error will be sent as a trailer to the client
|
||||
// ok will return nil, though.
|
||||
func writeStream(ctx context.Context, c *heartbeatconn.Conn, stream io.Reader, stype uint32) (errStream, errConn error) {
|
||||
debug("writeStream: enter stype=%v", stype)
|
||||
defer debug("writeStream: return")
|
||||
if stype == 0 {
|
||||
panic("stype must be non-zero")
|
||||
}
|
||||
if !IsPublicFrameType(stype) {
|
||||
panic(fmt.Sprintf("stype %v is not public", stype))
|
||||
}
|
||||
return doWriteStream(ctx, c, stream, stype)
|
||||
}
|
||||
|
||||
func doWriteStream(ctx context.Context, c *heartbeatconn.Conn, stream io.Reader, stype uint32) (errStream, errConn error) {
|
||||
|
||||
// RULE1 (buf == <zero>) XOR (err == nil)
|
||||
type read struct {
|
||||
buf base2bufpool.Buffer
|
||||
err error
|
||||
}
|
||||
|
||||
var wg sync.WaitGroup
|
||||
defer wg.Wait()
|
||||
reads := make(chan read, 5)
|
||||
var stopReading uint32
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer close(reads)
|
||||
for atomic.LoadUint32(&stopReading) == 0 {
|
||||
buffer := bufpool.Get(1 << FramePayloadShift)
|
||||
bufferBytes := buffer.Bytes()
|
||||
n, err := io.ReadFull(stream, bufferBytes)
|
||||
buffer.Shrink(uint(n))
|
||||
// if we received anything, send one read without an error (RULE 1)
|
||||
if n > 0 {
|
||||
reads <- read{buffer, nil}
|
||||
}
|
||||
if err == io.ErrUnexpectedEOF {
|
||||
// happens iff io.ReadFull read io.EOF from stream
|
||||
err = io.EOF
|
||||
}
|
||||
if err != nil {
|
||||
reads <- read{err: err} // RULE1
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
|
||||
defer func() {
|
||||
// stop reading
|
||||
atomic.StoreUint32(&stopReading, 1)
|
||||
// drain in-flight reads
|
||||
for read := range reads {
|
||||
debug("doWriteStream: drain read channel")
|
||||
read.buf.Free()
|
||||
}
|
||||
}()
|
||||
|
||||
for read := range reads {
|
||||
if read.err == nil {
|
||||
// RULE 1: read.buf is valid
|
||||
// next line is the hot path...
|
||||
writeErr := c.WriteFrame(read.buf.Bytes(), stype)
|
||||
read.buf.Free()
|
||||
if writeErr != nil {
|
||||
return nil, writeErr
|
||||
}
|
||||
continue
|
||||
} else if read.err == io.EOF {
|
||||
if err := c.WriteFrame([]byte{}, End); err != nil {
|
||||
return nil, err
|
||||
}
|
||||
break
|
||||
} else {
|
||||
errReader := strings.NewReader(read.err.Error())
|
||||
errReadErrReader, errConnWrite := doWriteStream(ctx, c, errReader, StreamErrTrailer)
|
||||
if errReadErrReader != nil {
|
||||
panic(errReadErrReader) // in-memory, cannot happen
|
||||
}
|
||||
return read.err, errConnWrite
|
||||
}
|
||||
}
|
||||
|
||||
return nil, nil
|
||||
}
|
||||
|
||||
type ReadStreamErrorKind int
|
||||
|
||||
const (
|
||||
ReadStreamErrorKindConn ReadStreamErrorKind = 1 + iota
|
||||
ReadStreamErrorKindWrite
|
||||
ReadStreamErrorKindSource
|
||||
ReadStreamErrorKindStreamErrTrailerEncoding
|
||||
ReadStreamErrorKindUnexpectedFrameType
|
||||
)
|
||||
|
||||
type ReadStreamError struct {
|
||||
Kind ReadStreamErrorKind
|
||||
Err error
|
||||
}
|
||||
|
||||
func (e *ReadStreamError) Error() string {
|
||||
kindStr := ""
|
||||
switch e.Kind {
|
||||
case ReadStreamErrorKindConn:
|
||||
kindStr = " read error: "
|
||||
case ReadStreamErrorKindWrite:
|
||||
kindStr = " write error: "
|
||||
case ReadStreamErrorKindSource:
|
||||
kindStr = " source error: "
|
||||
case ReadStreamErrorKindStreamErrTrailerEncoding:
|
||||
kindStr = " source implementation error: "
|
||||
case ReadStreamErrorKindUnexpectedFrameType:
|
||||
kindStr = " protocol error: "
|
||||
}
|
||||
return fmt.Sprintf("stream:%s%s", kindStr, e.Err)
|
||||
}
|
||||
|
||||
var _ net.Error = &ReadStreamError{}
|
||||
|
||||
func (e ReadStreamError) Timeout() bool {
|
||||
if netErr, ok := e.Err.(net.Error); ok {
|
||||
return netErr.Timeout()
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
// This function is deprecated in net.Error and since this
|
||||
// function is not involved in .Accept() code path, nothing
|
||||
// really needs this method to be here.
|
||||
func (e ReadStreamError) Temporary() bool {
|
||||
if te, ok := e.Err.(interface{ Temporary() bool }); ok {
|
||||
return te.Temporary()
|
||||
}
|
||||
return false
|
||||
}
|
||||
|
||||
var _ net.Error = &ReadStreamError{}
|
||||
|
||||
func (e ReadStreamError) IsReadError() bool {
|
||||
return e.Kind != ReadStreamErrorKindWrite
|
||||
}
|
||||
|
||||
func (e ReadStreamError) IsWriteError() bool {
|
||||
return e.Kind == ReadStreamErrorKindWrite
|
||||
}
|
||||
|
||||
type readFrameResult struct {
|
||||
f frameconn.Frame
|
||||
err error
|
||||
}
|
||||
|
||||
// readFrames reads from c into reads
|
||||
// if a read from c encounters an error, noMoreReads is closed before sending the result into reads
|
||||
func readFrames(reads chan<- readFrameResult, noMoreReads chan<- struct{}, c *heartbeatconn.Conn) {
|
||||
// noMoreReads is already closed, don't re-close it
|
||||
defer close(reads)
|
||||
for { // only exits after a read error, make sure noMoreReads is closed
|
||||
var r readFrameResult
|
||||
r.f, r.err = c.ReadFrame()
|
||||
if r.err != nil && noMoreReads != nil {
|
||||
close(noMoreReads)
|
||||
}
|
||||
reads <- r
|
||||
if r.err != nil {
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// ReadStream will close c if an error reading from c or writing to receiver occurs
|
||||
//
|
||||
// readStream calls itself recursively to read multi-frame error trailers
|
||||
// Thus, the reads channel needs to be a parameter.
|
||||
func readStream(reads <-chan readFrameResult, c *heartbeatconn.Conn, receiver io.Writer, stype uint32) *ReadStreamError {
|
||||
|
||||
var f frameconn.Frame
|
||||
for read := range reads {
|
||||
debug("readStream: read frame %v %v", read.f.Header, read.err)
|
||||
f = read.f
|
||||
if read.err != nil {
|
||||
return &ReadStreamError{ReadStreamErrorKindConn, read.err}
|
||||
}
|
||||
if f.Header.Type != stype {
|
||||
break
|
||||
}
|
||||
|
||||
n, err := receiver.Write(f.Buffer.Bytes())
|
||||
if err != nil {
|
||||
f.Buffer.Free()
|
||||
return &ReadStreamError{ReadStreamErrorKindWrite, err} // FIXME wrap as writer error
|
||||
}
|
||||
if n != len(f.Buffer.Bytes()) {
|
||||
f.Buffer.Free()
|
||||
return &ReadStreamError{ReadStreamErrorKindWrite, io.ErrShortWrite}
|
||||
}
|
||||
f.Buffer.Free()
|
||||
}
|
||||
|
||||
if f.Header.Type == End {
|
||||
debug("readStream: End reached")
|
||||
return nil
|
||||
}
|
||||
|
||||
if f.Header.Type == StreamErrTrailer {
|
||||
debug("readStream: begin of StreamErrTrailer")
|
||||
var errBuf bytes.Buffer
|
||||
if n, err := errBuf.Write(f.Buffer.Bytes()); n != len(f.Buffer.Bytes()) || err != nil {
|
||||
panic(fmt.Sprintf("unexpected bytes.Buffer write error: %v %v", n, err))
|
||||
}
|
||||
// recursion ftw! we won't enter this if stmt because stype == StreamErrTrailer in the following call
|
||||
rserr := readStream(reads, c, &errBuf, StreamErrTrailer)
|
||||
if rserr != nil && rserr.Kind == ReadStreamErrorKindWrite {
|
||||
panic(fmt.Sprintf("unexpected bytes.Buffer write error: %s", rserr))
|
||||
} else if rserr != nil {
|
||||
debug("readStream: rserr != nil && != ReadStreamErrorKindWrite: %v %v\n", rserr.Kind, rserr.Err)
|
||||
return rserr
|
||||
}
|
||||
if !utf8.Valid(errBuf.Bytes()) {
|
||||
return &ReadStreamError{ReadStreamErrorKindStreamErrTrailerEncoding, fmt.Errorf("source error, but not encoded as UTF-8")}
|
||||
}
|
||||
return &ReadStreamError{ReadStreamErrorKindSource, fmt.Errorf("%s", errBuf.String())}
|
||||
}
|
||||
|
||||
return &ReadStreamError{ReadStreamErrorKindUnexpectedFrameType, fmt.Errorf("unexpected frame type %v (expected %v)", f.Header.Type, stype)}
|
||||
}
|
||||
@@ -0,0 +1,296 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"io"
|
||||
"net"
|
||||
"sync"
|
||||
"sync/atomic"
|
||||
"time"
|
||||
|
||||
"github.com/pkg/errors"
|
||||
|
||||
"github.com/zrepl/zrepl/internal/rpc/dataconn/heartbeatconn"
|
||||
"github.com/zrepl/zrepl/internal/rpc/dataconn/timeoutconn"
|
||||
)
|
||||
|
||||
type Conn struct {
|
||||
hc *heartbeatconn.Conn
|
||||
|
||||
// whether the per-conn readFrames goroutine completed
|
||||
waitReadFramesDone chan struct{}
|
||||
// filled by per-conn readFrames goroutine
|
||||
frameReads chan readFrameResult
|
||||
|
||||
closeState closeState
|
||||
|
||||
// readMtx serializes read stream operations because we inherently only
|
||||
// support a single stream at a time over hc.
|
||||
readMtx sync.Mutex
|
||||
readClean bool
|
||||
|
||||
// writeMtx serializes write stream operations because we inherently only
|
||||
// support a single stream at a time over hc.
|
||||
writeMtx sync.Mutex
|
||||
writeClean bool
|
||||
}
|
||||
|
||||
var readMessageSentinel = fmt.Errorf("read stream complete")
|
||||
|
||||
var errWriteStreamToErrorUnknownState = fmt.Errorf("dataconn read stream: connection is in unknown state")
|
||||
|
||||
func Wrap(nc timeoutconn.Wire, sendHeartbeatInterval, peerTimeout time.Duration) *Conn {
|
||||
hc := heartbeatconn.Wrap(nc, sendHeartbeatInterval, peerTimeout)
|
||||
conn := &Conn{
|
||||
hc: hc, readClean: true, writeClean: true,
|
||||
waitReadFramesDone: make(chan struct{}),
|
||||
frameReads: make(chan readFrameResult, 5), // FIXME constant
|
||||
}
|
||||
go conn.readFrames()
|
||||
return conn
|
||||
}
|
||||
|
||||
func isConnCleanAfterRead(res *ReadStreamError) bool {
|
||||
return res == nil || res.Kind == ReadStreamErrorKindSource || res.Kind == ReadStreamErrorKindStreamErrTrailerEncoding
|
||||
}
|
||||
|
||||
func isConnCleanAfterWrite(err error) bool {
|
||||
return err == nil
|
||||
}
|
||||
|
||||
func (c *Conn) readFrames() {
|
||||
readFrames(c.frameReads, c.waitReadFramesDone, c.hc)
|
||||
}
|
||||
|
||||
func (c *Conn) ReadStreamedMessage(ctx context.Context, maxSize uint32, frameType uint32) (_ []byte, err *ReadStreamError) {
|
||||
|
||||
// if we are closed while reading, return that as an error
|
||||
if closeGuard, cse := c.closeState.RWEntry(); cse != nil {
|
||||
return nil, &ReadStreamError{
|
||||
Kind: ReadStreamErrorKindConn,
|
||||
Err: cse,
|
||||
}
|
||||
} else {
|
||||
defer func(err **ReadStreamError) {
|
||||
if closed := closeGuard.RWExit(); closed != nil {
|
||||
*err = &ReadStreamError{
|
||||
Kind: ReadStreamErrorKindConn,
|
||||
Err: closed,
|
||||
}
|
||||
}
|
||||
}(&err)
|
||||
}
|
||||
|
||||
c.readMtx.Lock()
|
||||
defer c.readMtx.Unlock()
|
||||
if !c.readClean {
|
||||
return nil, &ReadStreamError{
|
||||
Kind: ReadStreamErrorKindConn,
|
||||
Err: fmt.Errorf("dataconn read message: connection is in unknown state"),
|
||||
}
|
||||
}
|
||||
|
||||
r, w := io.Pipe()
|
||||
var buf bytes.Buffer
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
lr := io.LimitReader(r, int64(maxSize))
|
||||
if _, err := io.Copy(&buf, lr); err != nil && err != readMessageSentinel {
|
||||
panic(err)
|
||||
}
|
||||
}()
|
||||
err = readStream(c.frameReads, c.hc, w, frameType)
|
||||
c.readClean = isConnCleanAfterRead(err)
|
||||
_ = w.CloseWithError(readMessageSentinel) // always returns nil
|
||||
wg.Wait()
|
||||
if err != nil {
|
||||
return nil, err
|
||||
} else {
|
||||
return buf.Bytes(), nil
|
||||
}
|
||||
}
|
||||
|
||||
type StreamReader struct {
|
||||
*io.PipeReader
|
||||
conn *Conn
|
||||
closeConnOnClose bool
|
||||
}
|
||||
|
||||
func (r *StreamReader) Close() error {
|
||||
err := r.PipeReader.Close()
|
||||
if r.closeConnOnClose {
|
||||
r.conn.Close() // TODO error logging
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
// WriteStreamTo reads a stream from Conn and writes it to w.
|
||||
func (c *Conn) ReadStream(frameType uint32, closeConnOnClose bool) (_ *StreamReader, err error) {
|
||||
|
||||
// if we are closed while writing, return that as an error
|
||||
if closeGuard, cse := c.closeState.RWEntry(); cse != nil {
|
||||
return nil, cse
|
||||
} else {
|
||||
defer func(err *error) {
|
||||
if closed := closeGuard.RWExit(); closed != nil {
|
||||
*err = closed
|
||||
}
|
||||
}(&err)
|
||||
}
|
||||
|
||||
c.readMtx.Lock()
|
||||
if !c.readClean {
|
||||
c.readMtx.Unlock()
|
||||
return nil, errWriteStreamToErrorUnknownState
|
||||
}
|
||||
|
||||
r, w := io.Pipe()
|
||||
go func() {
|
||||
defer c.readMtx.Unlock()
|
||||
var err *ReadStreamError = readStream(c.frameReads, c.hc, w, frameType)
|
||||
if err != nil {
|
||||
_ = w.CloseWithError(err) // doc guarantees that error will always be nil
|
||||
} else {
|
||||
w.Close()
|
||||
}
|
||||
c.readClean = isConnCleanAfterRead(err)
|
||||
}()
|
||||
|
||||
return &StreamReader{PipeReader: r, conn: c, closeConnOnClose: closeConnOnClose}, nil
|
||||
}
|
||||
|
||||
func (c *Conn) WriteStreamedMessage(ctx context.Context, buf io.Reader, frameType uint32) (err error) {
|
||||
|
||||
// if we are closed while writing, return that as an error
|
||||
if closeGuard, cse := c.closeState.RWEntry(); cse != nil {
|
||||
return cse
|
||||
} else {
|
||||
defer func(err *error) {
|
||||
if closed := closeGuard.RWExit(); closed != nil {
|
||||
*err = closed
|
||||
}
|
||||
}(&err)
|
||||
}
|
||||
|
||||
c.writeMtx.Lock()
|
||||
defer c.writeMtx.Unlock()
|
||||
if !c.writeClean {
|
||||
return fmt.Errorf("dataconn write message: connection is in unknown state")
|
||||
}
|
||||
errBuf, errConn := writeStream(ctx, c.hc, buf, frameType)
|
||||
if errBuf != nil {
|
||||
panic(errBuf)
|
||||
}
|
||||
c.writeClean = isConnCleanAfterWrite(errConn)
|
||||
return errConn
|
||||
}
|
||||
|
||||
func (c *Conn) SendStream(ctx context.Context, stream io.ReadCloser, frameType uint32) (err error) {
|
||||
|
||||
// if we are closed while reading, return that as an error
|
||||
if closeGuard, cse := c.closeState.RWEntry(); cse != nil {
|
||||
return cse
|
||||
} else {
|
||||
defer func(err *error) {
|
||||
if closed := closeGuard.RWExit(); closed != nil {
|
||||
*err = closed
|
||||
}
|
||||
}(&err)
|
||||
}
|
||||
|
||||
c.writeMtx.Lock()
|
||||
defer c.writeMtx.Unlock()
|
||||
if !c.writeClean {
|
||||
return fmt.Errorf("dataconn send stream: connection is in unknown state")
|
||||
}
|
||||
|
||||
errStream, errConn := writeStream(ctx, c.hc, stream, frameType)
|
||||
|
||||
c.writeClean = isConnCleanAfterWrite(errConn) // TODO correct?
|
||||
|
||||
if errStream != nil {
|
||||
return errStream
|
||||
} else if errConn != nil {
|
||||
return errConn
|
||||
}
|
||||
// TODO combined error?
|
||||
return nil
|
||||
}
|
||||
|
||||
type closeState struct {
|
||||
closeCount uint32
|
||||
}
|
||||
|
||||
type closeStateErrConnectionClosed struct{}
|
||||
|
||||
var _ net.Error = (*closeStateErrConnectionClosed)(nil)
|
||||
|
||||
func (e *closeStateErrConnectionClosed) Error() string {
|
||||
return "connection closed"
|
||||
}
|
||||
|
||||
func (e *closeStateErrConnectionClosed) Timeout() bool { return false }
|
||||
|
||||
// This function is deprecated in net.Error and since this
|
||||
// function is not involved in .Accept() code path, nothing
|
||||
// really needs this method to be here.
|
||||
func (e *closeStateErrConnectionClosed) Temporary() bool { return false }
|
||||
|
||||
func (s *closeState) CloseEntry() error {
|
||||
firstCloser := atomic.AddUint32(&s.closeCount, 1) == 1
|
||||
if !firstCloser {
|
||||
return errors.New("duplicate close")
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
type closeStateEntry struct {
|
||||
s *closeState
|
||||
entryCount uint32
|
||||
}
|
||||
|
||||
func (s *closeState) RWEntry() (e *closeStateEntry, err net.Error) {
|
||||
entry := &closeStateEntry{s, atomic.LoadUint32(&s.closeCount)}
|
||||
if entry.entryCount > 0 {
|
||||
return nil, &closeStateErrConnectionClosed{}
|
||||
}
|
||||
return entry, nil
|
||||
}
|
||||
|
||||
func (e *closeStateEntry) RWExit() net.Error {
|
||||
if atomic.LoadUint32(&e.entryCount) == e.entryCount {
|
||||
// no calls to Close() while running rw operation
|
||||
return nil
|
||||
}
|
||||
return &closeStateErrConnectionClosed{}
|
||||
}
|
||||
|
||||
func (c *Conn) Close() error {
|
||||
if err := c.closeState.CloseEntry(); err != nil {
|
||||
return errors.Wrap(err, "stream conn close")
|
||||
}
|
||||
|
||||
// Shutdown c.hc, which will cause c.readFrames to close c.waitReadFramesDone
|
||||
err := c.hc.Shutdown()
|
||||
// However, c.readFrames may be blocking on a filled c.frameReads
|
||||
// and since the connection is closed, nobody is going to read from it
|
||||
for read := range c.frameReads {
|
||||
debug("Conn.Close() draining queued read")
|
||||
read.f.Buffer.Free()
|
||||
}
|
||||
// if we can't close, don't expect c.readFrames to terminate
|
||||
// this might leak c.readFrames, but we can't do something useful at this point
|
||||
if err != nil {
|
||||
return err
|
||||
}
|
||||
// shutdown didn't report errors, so c.readFrames will exit due to read error
|
||||
// This behavior is the contract we have with c.hc.
|
||||
// If that contract is broken, this read will block indefinitely and
|
||||
// cause an easily diagnosable goroutine leak (of this goroutine)
|
||||
<-c.waitReadFramesDone
|
||||
return nil
|
||||
}
|
||||
@@ -0,0 +1,21 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"os"
|
||||
)
|
||||
|
||||
var debugEnabled bool = false
|
||||
|
||||
func init() {
|
||||
if os.Getenv("ZREPL_RPC_DATACONN_STREAM_DEBUG") != "" {
|
||||
debugEnabled = true
|
||||
}
|
||||
}
|
||||
|
||||
//nolint:deadcode,unused
|
||||
func debug(format string, args ...interface{}) {
|
||||
if debugEnabled {
|
||||
fmt.Fprintf(os.Stderr, "rpc/dataconn/stream: %s\n", fmt.Sprintf(format, args...))
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,132 @@
|
||||
package stream
|
||||
|
||||
import (
|
||||
"bytes"
|
||||
"context"
|
||||
"fmt"
|
||||
"strings"
|
||||
"sync"
|
||||
"testing"
|
||||
"time"
|
||||
|
||||
"github.com/stretchr/testify/assert"
|
||||
"github.com/stretchr/testify/require"
|
||||
|
||||
"github.com/zrepl/zrepl/internal/logger"
|
||||
"github.com/zrepl/zrepl/internal/rpc/dataconn/heartbeatconn"
|
||||
"github.com/zrepl/zrepl/internal/util/socketpair"
|
||||
)
|
||||
|
||||
func TestFrameTypesOk(t *testing.T) {
|
||||
t.Logf("%v", End)
|
||||
assert.True(t, heartbeatconn.IsPublicFrameType(End))
|
||||
assert.True(t, heartbeatconn.IsPublicFrameType(StreamErrTrailer))
|
||||
}
|
||||
|
||||
func TestStreamer(t *testing.T) {
|
||||
|
||||
anc, bnc, err := socketpair.SocketPair()
|
||||
require.NoError(t, err)
|
||||
|
||||
hto := 1 * time.Hour
|
||||
a := heartbeatconn.Wrap(anc, hto, hto)
|
||||
b := heartbeatconn.Wrap(bnc, hto, hto)
|
||||
|
||||
log := logger.NewStderrDebugLogger()
|
||||
ctx := WithLogger(context.Background(), log)
|
||||
|
||||
stype := uint32(0x23)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
var buf bytes.Buffer
|
||||
buf.Write(
|
||||
bytes.Repeat([]byte{1, 2}, 1<<25),
|
||||
)
|
||||
writeStream(ctx, a, &buf, stype)
|
||||
log.Debug("WriteStream returned")
|
||||
a.Shutdown()
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
var buf bytes.Buffer
|
||||
ch := make(chan readFrameResult, 5)
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
readFrames(ch, nil, b)
|
||||
}()
|
||||
err := readStream(ch, b, &buf, stype)
|
||||
log.WithField("errType", fmt.Sprintf("%T %v", err, err)).Debug("ReadStream returned")
|
||||
assert.Nil(t, err)
|
||||
expected := bytes.Repeat([]byte{1, 2}, 1<<25)
|
||||
assert.True(t, bytes.Equal(expected, buf.Bytes()))
|
||||
b.Shutdown()
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
|
||||
}
|
||||
|
||||
type errReader struct {
|
||||
t *testing.T
|
||||
readErr error
|
||||
}
|
||||
|
||||
func (er errReader) Read(p []byte) (n int, err error) {
|
||||
er.t.Logf("errReader.Read called")
|
||||
return 0, er.readErr
|
||||
}
|
||||
|
||||
func TestMultiFrameStreamErrTraileror(t *testing.T) {
|
||||
anc, bnc, err := socketpair.SocketPair()
|
||||
require.NoError(t, err)
|
||||
|
||||
hto := 1 * time.Hour
|
||||
a := heartbeatconn.Wrap(anc, hto, hto)
|
||||
b := heartbeatconn.Wrap(bnc, hto, hto)
|
||||
|
||||
log := logger.NewStderrDebugLogger()
|
||||
ctx := WithLogger(context.Background(), log)
|
||||
|
||||
longErr := fmt.Errorf("an error that definitley spans more than one frame:\n%s", strings.Repeat("a\n", 1<<4))
|
||||
|
||||
stype := uint32(0x23)
|
||||
|
||||
var wg sync.WaitGroup
|
||||
wg.Add(2)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
r := errReader{t, longErr}
|
||||
writeStream(ctx, a, &r, stype)
|
||||
a.Shutdown()
|
||||
}()
|
||||
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
defer b.Shutdown()
|
||||
var buf bytes.Buffer
|
||||
ch := make(chan readFrameResult, 5)
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
readFrames(ch, nil, b)
|
||||
}()
|
||||
err := readStream(ch, b, &buf, stype)
|
||||
t.Logf("%s", err)
|
||||
require.NotNil(t, err)
|
||||
assert.True(t, buf.Len() == 0)
|
||||
assert.Equal(t, err.Kind, ReadStreamErrorKindSource)
|
||||
receivedErr := err.Err.Error()
|
||||
expectedErr := longErr.Error()
|
||||
assert.True(t, receivedErr == expectedErr) // builtin Equals is too slow
|
||||
if receivedErr != expectedErr {
|
||||
t.Logf("lengths: %v %v", len(receivedErr), len(expectedErr))
|
||||
}
|
||||
}()
|
||||
|
||||
wg.Wait()
|
||||
}
|
||||
Reference in New Issue
Block a user