move implementation to internal/ directory (#828)

This commit is contained in:
Christian Schwarz
2024-10-18 19:21:17 +02:00
committed by GitHub
parent b9b9ad10cf
commit 908807bd59
360 changed files with 507 additions and 507 deletions
@@ -0,0 +1,169 @@
package base2bufpool
import (
"fmt"
"math/bits"
"sync"
)
type pool struct {
mtx sync.Mutex
bufs [][]byte
shift uint
}
func (p *pool) Put(buf []byte) {
p.mtx.Lock()
defer p.mtx.Unlock()
if len(buf) != 1<<p.shift {
panic(fmt.Sprintf("implementation error: %v %v", len(buf), 1<<p.shift))
}
if len(p.bufs) > 10 { // FIXME constant
return
}
p.bufs = append(p.bufs, buf)
}
func (p *pool) Get() []byte {
p.mtx.Lock()
defer p.mtx.Unlock()
if len(p.bufs) > 0 {
ret := p.bufs[len(p.bufs)-1]
p.bufs = p.bufs[0 : len(p.bufs)-1]
return ret
}
return make([]byte, 1<<p.shift)
}
type Pool struct {
minShift uint
maxShift uint
pools []pool
onNoFit NoFitBehavior
}
type Buffer struct {
// always power of 2, from Pool.pools
shiftBuf []byte
// presentedLen
payloadLen uint
// backref too pool for Free
pool *Pool
}
func (b Buffer) Bytes() []byte {
return b.shiftBuf[0:b.payloadLen]
}
func (b *Buffer) Shrink(newPayloadLen uint) {
if newPayloadLen > b.payloadLen {
panic(fmt.Sprintf("shrink is actually an expand, invalid: %v %v", newPayloadLen, b.payloadLen))
}
b.payloadLen = newPayloadLen
}
func (b Buffer) Free() {
if b.pool != nil {
b.pool.put(b)
}
}
//go:generate enumer -type NoFitBehavior
type NoFitBehavior uint
const (
AllocateSmaller NoFitBehavior = 1 << iota
AllocateLarger
Allocate NoFitBehavior = AllocateSmaller | AllocateLarger
Panic NoFitBehavior = 0
)
func New(minShift, maxShift uint, noFitBehavior NoFitBehavior) *Pool {
if minShift > 63 || maxShift > 63 {
panic(fmt.Sprintf("{min|max}Shift are the _exponent_, got minShift=%v maxShift=%v and limit of 63, which amounts to %v bits", minShift, maxShift, uint64(1)<<63))
}
pools := make([]pool, maxShift-minShift+1)
for i := uint(0); i < uint(len(pools)); i++ {
i := i // the closure below must copy i
pools[i] = pool{
shift: minShift + i,
bufs: make([][]byte, 0, 10),
}
}
return &Pool{
minShift: minShift,
maxShift: maxShift,
pools: pools,
onNoFit: noFitBehavior,
}
}
func fittingShift(x uint) uint {
if x == 0 {
return 0
}
blen := uint(bits.Len(x))
if 1<<(blen-1) == x {
return blen - 1
}
return blen
}
func (p *Pool) handlePotentialNoFit(reqShift uint) (buf Buffer, didHandle bool) {
if reqShift == 0 {
if p.onNoFit&AllocateSmaller != 0 {
return Buffer{[]byte{}, 0, nil}, true
} else {
goto doPanic
}
}
if reqShift < p.minShift {
if p.onNoFit&AllocateSmaller != 0 {
goto alloc
} else {
goto doPanic
}
}
if reqShift > p.maxShift {
if p.onNoFit&AllocateLarger != 0 {
goto alloc
} else {
goto doPanic
}
}
return Buffer{}, false
alloc:
return Buffer{make([]byte, 1<<reqShift), 1 << reqShift, nil}, true
doPanic:
panic(fmt.Sprintf("base2bufpool: configured to panic on shift=%v (minShift=%v maxShift=%v)", reqShift, p.minShift, p.maxShift))
}
func (p *Pool) Get(minSize uint) Buffer {
shift := fittingShift(minSize)
buf, didHandle := p.handlePotentialNoFit(shift)
if didHandle {
buf.Shrink(minSize)
return buf
}
idx := int64(shift) - int64(p.minShift)
return Buffer{p.pools[idx].Get(), minSize, p}
}
func (p *Pool) put(buffer Buffer) {
if buffer.pool != p {
panic("putting buffer to pool where it didn't originate from")
}
buf := buffer.shiftBuf
if bits.OnesCount(uint(len(buf))) > 1 {
panic(fmt.Sprintf("putting buffer that is not power of two len: %v", len(buf)))
}
if len(buf) == 0 {
return
}
shift := fittingShift(uint(len(buf)))
if shift < p.minShift || shift > p.maxShift {
return // drop it
}
p.pools[shift-p.minShift].Put(buf)
}
@@ -0,0 +1,98 @@
package base2bufpool
import (
"fmt"
"testing"
"github.com/stretchr/testify/assert"
)
func TestPoolAllocBehavior(t *testing.T) {
type testcase struct {
poolMinShift, poolMaxShift uint
behavior NoFitBehavior
get uint
expShiftBufLen int64 // -1 if panic expected
}
tcs := []testcase{
{
15, 20, Allocate,
1 << 14, 1 << 14,
},
{
15, 20, Allocate,
1 << 22, 1 << 22,
},
{
15, 20, Panic,
1 << 16, 1 << 16,
},
{
15, 20, Panic,
1 << 14, -1,
},
{
15, 20, Panic,
1 << 22, -1,
},
{
15, 20, Panic,
(1 << 15) + 23, 1 << 16,
},
{
15, 20, Panic,
0, -1, // yep, 0 always works, even
},
{
15, 20, Allocate,
0, 0,
},
{
15, 20, AllocateSmaller,
1 << 14, 1 << 14,
},
{
15, 20, AllocateSmaller,
1 << 22, -1,
},
}
for i := range tcs {
tc := tcs[i]
t.Run(fmt.Sprintf("[%d,%d] behav=%s Get(%d) exp=%d", tc.poolMinShift, tc.poolMaxShift, tc.behavior, tc.get, tc.expShiftBufLen), func(t *testing.T) {
pool := New(tc.poolMinShift, tc.poolMaxShift, tc.behavior)
if tc.expShiftBufLen == -1 {
assert.Panics(t, func() {
pool.Get(tc.get)
})
return
}
buf := pool.Get(tc.get)
assert.True(t, uint(len(buf.Bytes())) == tc.get)
assert.True(t, int64(len(buf.shiftBuf)) == tc.expShiftBufLen)
})
}
}
func TestFittingShift(t *testing.T) {
assert.Equal(t, uint(16), fittingShift(1+1<<15))
assert.Equal(t, uint(15), fittingShift(1<<15))
}
func TestFreeFromPoolRangeDoesNotPanic(t *testing.T) {
pool := New(15, 20, Allocate)
buf := pool.Get(1 << 16)
assert.NotPanics(t, func() {
buf.Free()
})
}
func TestFreeFromOutOfPoolRangeDoesNotPanic(t *testing.T) {
pool := New(15, 20, Allocate)
buf := pool.Get(1 << 23)
assert.NotPanics(t, func() {
buf.Free()
})
}
@@ -0,0 +1,51 @@
// Code generated by "enumer -type NoFitBehavior"; DO NOT EDIT.
package base2bufpool
import (
"fmt"
)
const _NoFitBehaviorName = "PanicAllocateSmallerAllocateLargerAllocate"
var _NoFitBehaviorIndex = [...]uint8{0, 5, 20, 34, 42}
func (i NoFitBehavior) String() string {
if i >= NoFitBehavior(len(_NoFitBehaviorIndex)-1) {
return fmt.Sprintf("NoFitBehavior(%d)", i)
}
return _NoFitBehaviorName[_NoFitBehaviorIndex[i]:_NoFitBehaviorIndex[i+1]]
}
var _NoFitBehaviorValues = []NoFitBehavior{0, 1, 2, 3}
var _NoFitBehaviorNameToValueMap = map[string]NoFitBehavior{
_NoFitBehaviorName[0:5]: 0,
_NoFitBehaviorName[5:20]: 1,
_NoFitBehaviorName[20:34]: 2,
_NoFitBehaviorName[34:42]: 3,
}
// NoFitBehaviorString retrieves an enum value from the enum constants string name.
// Throws an error if the param is not part of the enum.
func NoFitBehaviorString(s string) (NoFitBehavior, error) {
if val, ok := _NoFitBehaviorNameToValueMap[s]; ok {
return val, nil
}
return 0, fmt.Errorf("%s does not belong to NoFitBehavior values", s)
}
// NoFitBehaviorValues returns all values of the enum
func NoFitBehaviorValues() []NoFitBehavior {
return _NoFitBehaviorValues
}
// IsANoFitBehavior returns "true" if the value is listed in the enum definition. "false" otherwise
func (i NoFitBehavior) IsANoFitBehavior() bool {
for _, v := range _NoFitBehaviorValues {
if i == v {
return true
}
}
return false
}
+227
View File
@@ -0,0 +1,227 @@
package dataconn
import (
"bytes"
"context"
"fmt"
"io"
"strings"
"google.golang.org/protobuf/proto"
"github.com/zrepl/zrepl/internal/replication/logic/pdu"
"github.com/zrepl/zrepl/internal/rpc/dataconn/stream"
"github.com/zrepl/zrepl/internal/transport"
)
type Client struct {
log Logger
cn transport.Connecter
}
func NewClient(connecter transport.Connecter, log Logger) *Client {
return &Client{
log: log,
cn: connecter,
}
}
func (c *Client) send(ctx context.Context, conn *stream.Conn, endpoint string, req proto.Message, stream io.ReadCloser) error {
var buf bytes.Buffer
_, memErr := buf.WriteString(endpoint)
if memErr != nil {
panic(memErr)
}
if err := conn.WriteStreamedMessage(ctx, &buf, ReqHeader); err != nil {
return err
}
protobufBytes, err := proto.Marshal(req)
if err != nil {
return err
}
protobuf := bytes.NewBuffer(protobufBytes)
if err := conn.WriteStreamedMessage(ctx, protobuf, ReqStructured); err != nil {
return err
}
if stream != nil {
return conn.SendStream(ctx, stream, ZFSStream)
} else {
return nil
}
}
type RemoteHandlerError struct {
msg string
}
func (e *RemoteHandlerError) Error() string {
return fmt.Sprintf("server error: %s", e.msg)
}
type ProtocolError struct {
cause error
}
func (e *ProtocolError) Error() string {
return fmt.Sprintf("protocol error: %s", e.cause)
}
func (c *Client) recv(ctx context.Context, conn *stream.Conn, res proto.Message) error {
headerBuf, err := conn.ReadStreamedMessage(ctx, ResponseHeaderMaxSize, ResHeader)
if err != nil {
return err
}
header := string(headerBuf)
if strings.HasPrefix(header, responseHeaderHandlerErrorPrefix) {
// FIXME distinguishable error type
return &RemoteHandlerError{strings.TrimPrefix(header, responseHeaderHandlerErrorPrefix)}
}
if !strings.HasPrefix(header, responseHeaderHandlerOk) {
return &ProtocolError{fmt.Errorf("invalid header: %q", header)}
}
protobuf, err := conn.ReadStreamedMessage(ctx, ResponseStructuredMaxSize, ResStructured)
if err != nil {
return err
}
if err := proto.Unmarshal(protobuf, res); err != nil {
return &ProtocolError{fmt.Errorf("cannot unmarshal structured part of response: %s", err)}
}
return nil
}
func (c *Client) getWire(ctx context.Context) (*stream.Conn, error) {
nc, err := c.cn.Connect(ctx)
if err != nil {
return nil, err
}
conn := stream.Wrap(nc, HeartbeatInterval, HeartbeatPeerTimeout)
return conn, nil
}
func (c *Client) putWire(conn *stream.Conn) {
if err := conn.Close(); err != nil {
c.log.WithError(err).Error("error closing connection")
}
}
func (c *Client) ReqSend(ctx context.Context, req *pdu.SendReq) (*pdu.SendRes, io.ReadCloser, error) {
conn, err := c.getWire(ctx)
if err != nil {
return nil, nil, err
}
if err := c.send(ctx, conn, EndpointSend, req, nil); err != nil {
return nil, nil, err
}
var res pdu.SendRes
if err := c.recv(ctx, conn, &res); err != nil {
return nil, nil, err
}
var stream io.ReadCloser
stream, err = conn.ReadStream(ZFSStream, true) // no shadow
if err != nil {
return nil, nil, err
}
return &res, stream, nil
}
func (c *Client) ReqRecv(ctx context.Context, req *pdu.ReceiveReq, stream io.ReadCloser) (*pdu.ReceiveRes, error) {
defer c.log.Debug("ReqRecv returns")
conn, err := c.getWire(ctx)
if err != nil {
return nil, err
}
// send and recv response concurrently to catch early exists of remote handler
// (e.g. disk full, permission error, etc)
type recvRes struct {
res *pdu.ReceiveRes
err error
}
recvErrChan := make(chan recvRes)
go func() {
res := &pdu.ReceiveRes{}
if err := c.recv(ctx, conn, res); err != nil {
recvErrChan <- recvRes{res, err}
} else {
recvErrChan <- recvRes{res, nil}
}
}()
sendErrChan := make(chan error)
go func() {
if err := c.send(ctx, conn, EndpointRecv, req, stream); err != nil {
sendErrChan <- err
} else {
sendErrChan <- nil
}
}()
var res recvRes
var sendErr error
var cause error // one of the above
didTryClose := false
for i := 0; i < 2; i++ {
select {
case res = <-recvErrChan:
c.log.WithField("errType", fmt.Sprintf("%T", res.err)).WithError(res.err).Debug("recv goroutine returned")
if res.err != nil && cause == nil {
cause = res.err
}
case sendErr = <-sendErrChan:
c.log.WithField("errType", fmt.Sprintf("%T", sendErr)).WithError(sendErr).Debug("send goroutine returned")
if sendErr != nil && cause == nil {
cause = sendErr
}
}
if !didTryClose && (res.err != nil || sendErr != nil) {
didTryClose = true
if err := conn.Close(); err != nil {
c.log.WithError(err).Error("ReqRecv: cannot close connection, will likely block indefinitely")
}
c.log.WithError(err).Debug("ReqRecv: closed connection, should trigger other goroutine error")
}
}
if !didTryClose {
// didn't close it in above loop, so we can give it back
c.putWire(conn)
}
// if receive failed with a RemoteHandlerError, we know the transport was not broken
// => take the remote error as cause for the operation to fail
// TODO combine errors if send also failed
// (after all, send could have crashed on our side, rendering res.err a mere symptom of the cause)
if _, ok := res.err.(*RemoteHandlerError); ok {
cause = res.err
}
return res.res, cause
}
func (c *Client) ReqPing(ctx context.Context, req *pdu.PingReq) (*pdu.PingRes, error) {
conn, err := c.getWire(ctx)
if err != nil {
return nil, err
}
defer c.putWire(conn)
if err := c.send(ctx, conn, EndpointPing, req, nil); err != nil {
return nil, err
}
var res pdu.PingRes
if err := c.recv(ctx, conn, &res); err != nil {
return nil, err
}
return &res, nil
}
+22
View File
@@ -0,0 +1,22 @@
package dataconn
import (
"fmt"
"os"
)
//nolint:unused
var debugEnabled bool = false
func init() {
if os.Getenv("ZREPL_RPC_DATACONN_DEBUG") != "" {
debugEnabled = true
}
}
//nolint:deadcode,unused
func debug(format string, args ...interface{}) {
if debugEnabled {
fmt.Fprintf(os.Stderr, "rpc/dataconn: %s\n", fmt.Sprintf(format, args...))
}
}
+256
View File
@@ -0,0 +1,256 @@
package dataconn
import (
"bytes"
"context"
"fmt"
"io"
"sync"
"google.golang.org/protobuf/proto"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/replication/logic/pdu"
"github.com/zrepl/zrepl/internal/rpc/dataconn/stream"
"github.com/zrepl/zrepl/internal/transport"
)
// WireInterceptor has a chance to exchange the context and connection on each client connection.
type WireInterceptor func(ctx context.Context, rawConn *transport.AuthConn) (context.Context, *transport.AuthConn)
// Handler implements the functionality that is exposed by Server to the Client.
type Handler interface {
// Send handles a SendRequest.
// The returned io.ReadCloser is allowed to be nil, for example if the requested Send is a dry-run.
Send(ctx context.Context, r *pdu.SendReq) (*pdu.SendRes, io.ReadCloser, error)
// Receive handles a ReceiveRequest.
// It is guaranteed that Server calls Receive with a stream that holds the IdleConnTimeout
// configured in ServerConfig.Shared.IdleConnTimeout.
Receive(ctx context.Context, r *pdu.ReceiveReq, receive io.ReadCloser) (*pdu.ReceiveRes, error)
// PingDataconn handles a PingReq
PingDataconn(ctx context.Context, r *pdu.PingReq) (*pdu.PingRes, error)
}
type Logger = logger.Logger
type ContextInterceptorData interface {
FullMethod() string
ClientIdentity() string
}
type ContextInterceptor = func(ctx context.Context, data ContextInterceptorData, handler func(ctx context.Context))
type Server struct {
h Handler
wi WireInterceptor
ci ContextInterceptor
log Logger
}
var noopContextInteceptor = func(ctx context.Context, _ ContextInterceptorData, handler func(context.Context)) {
handler(ctx)
}
// wi and ci may be nil
func NewServer(wi WireInterceptor, ci ContextInterceptor, logger Logger, handler Handler) *Server {
if ci == nil {
ci = noopContextInteceptor
}
return &Server{
h: handler,
wi: wi,
ci: ci,
log: logger,
}
}
// Serve consumes the listener, closes it as soon as ctx is closed.
// No accept errors are returned: they are logged to the Logger passed
// to the constructor.
func (s *Server) Serve(ctx context.Context, l transport.AuthenticatedListener) {
var wg sync.WaitGroup
defer wg.Wait()
ctx, cancel := context.WithCancel(ctx)
defer cancel()
wg.Add(1)
go func() {
defer wg.Done()
<-ctx.Done()
s.log.Debug("context done, closing listener")
if err := l.Close(); err != nil {
s.log.WithError(err).Error("cannot close listener")
}
}()
conns := make(chan *transport.AuthConn)
wg.Add(1)
go func() {
defer wg.Done()
defer close(conns)
for {
conn, err := l.Accept(ctx)
if err != nil {
if ctx.Done() != nil {
s.log.Debug("stop accepting after context is done")
return
}
s.log.WithError(err).Error("accept error")
continue
}
conns <- conn
}
}()
for conn := range conns {
wg.Add(1)
go func(conn *transport.AuthConn) {
defer wg.Done()
s.serveConn(conn)
}(conn)
}
}
type contextInterceptorData struct {
fullMethod string
clientIdentity string
}
func (d contextInterceptorData) FullMethod() string { return d.fullMethod }
func (d contextInterceptorData) ClientIdentity() string { return d.clientIdentity }
func (s *Server) serveConn(nc *transport.AuthConn) {
s.log.Debug("serveConn begin")
defer s.log.Debug("serveConn done")
ctx := context.Background()
if s.wi != nil {
ctx, nc = s.wi(ctx, nc)
}
c := stream.Wrap(nc, HeartbeatInterval, HeartbeatPeerTimeout)
defer func() {
s.log.Debug("close client connection")
if err := c.Close(); err != nil {
s.log.WithError(err).Error("cannot close client connection")
}
}()
header, err := c.ReadStreamedMessage(ctx, RequestHeaderMaxSize, ReqHeader)
if err != nil {
s.log.WithError(err).Error("error reading structured part")
return
}
endpoint := string(header)
data := contextInterceptorData{
fullMethod: endpoint,
clientIdentity: nc.ClientIdentity(),
}
s.ci(ctx, data, func(ctx context.Context) {
s.serveConnRequest(ctx, endpoint, c)
})
}
func (s *Server) serveConnRequest(ctx context.Context, endpoint string, c *stream.Conn) {
reqStructured, err := c.ReadStreamedMessage(ctx, RequestStructuredMaxSize, ReqStructured)
if err != nil {
s.log.WithError(err).Error("error reading structured part")
return
}
s.log.WithField("endpoint", endpoint).Debug("calling handler")
var res proto.Message
var sendStream io.ReadCloser
var handlerErr error
switch endpoint {
case EndpointSend:
var req pdu.SendReq
if err := proto.Unmarshal(reqStructured, &req); err != nil {
s.log.WithError(err).Error("cannot unmarshal send request")
return
}
res, sendStream, handlerErr = s.h.Send(ctx, &req) // SHADOWING
// ensure that we always close the sendStream
if sendStream != nil {
defer func() {
err := sendStream.Close()
if err != nil {
s.log.WithError(err).Error("cannot close send stream")
}
}()
}
case EndpointRecv:
var req pdu.ReceiveReq
if err := proto.Unmarshal(reqStructured, &req); err != nil {
s.log.WithError(err).Error("cannot unmarshal receive request")
return
}
stream, err := c.ReadStream(ZFSStream, false)
if err != nil {
s.log.WithError(err).Error("cannot open stream in receive request")
return
}
res, handlerErr = s.h.Receive(ctx, &req, stream) // SHADOWING
case EndpointPing:
var req pdu.PingReq
if err := proto.Unmarshal(reqStructured, &req); err != nil {
s.log.WithError(err).Error("cannot unmarshal ping request")
return
}
res, handlerErr = s.h.PingDataconn(ctx, &req) // SHADOWING
default:
s.log.WithField("endpoint", endpoint).Error("unknown endpoint")
handlerErr = fmt.Errorf("requested endpoint does not exist")
}
s.log.WithField("endpoint", endpoint).WithField("errType", fmt.Sprintf("%T", handlerErr)).Debug("handler returned")
// prepare protobuf now to return the protobuf error in the header
// if marshaling fails. We consider failed marshaling a handler error
var protobuf *bytes.Buffer
if handlerErr == nil {
if res == nil {
handlerErr = fmt.Errorf("implementation error: handler for endpoint %q returns nil error and nil result", endpoint)
s.log.WithError(err).Error("handle implementation error")
} else {
protobufBytes, err := proto.Marshal(res)
if err != nil {
s.log.WithError(err).Error("cannot marshal handler protobuf")
handlerErr = err
}
protobuf = bytes.NewBuffer(protobufBytes) // SHADOWING
}
}
var resHeaderBuf bytes.Buffer
if handlerErr == nil {
resHeaderBuf.WriteString(responseHeaderHandlerOk)
} else {
resHeaderBuf.WriteString(responseHeaderHandlerErrorPrefix)
resHeaderBuf.WriteString(handlerErr.Error())
}
if err := c.WriteStreamedMessage(ctx, &resHeaderBuf, ResHeader); err != nil {
s.log.WithError(err).Error("cannot write response header")
return
}
if handlerErr != nil {
s.log.Debug("early exit after handler error")
return
}
if err := c.WriteStreamedMessage(ctx, protobuf, ResStructured); err != nil {
s.log.WithError(err).Error("cannot write structured part of response")
return
}
if sendStream != nil {
err := c.SendStream(ctx, sendStream, ZFSStream)
if err != nil {
s.log.WithError(err).Error("cannot write send stream")
}
// sendStream.Close() done via defer above
}
}
+36
View File
@@ -0,0 +1,36 @@
package dataconn
import (
"time"
)
const (
EndpointPing string = "/v1/ping"
EndpointSend string = "/v1/send"
EndpointRecv string = "/v1/recv"
)
const (
ReqHeader uint32 = 1 + iota
ReqStructured
ResHeader
ResStructured
ZFSStream
)
// Note that changing theses constants may break interop with other clients
// Aggressive with timing, conservative (future compatible) with buffer sizes
const (
HeartbeatInterval = 5 * time.Second
HeartbeatPeerTimeout = 10 * time.Second
RequestHeaderMaxSize = 1 << 15
RequestStructuredMaxSize = 1 << 22
ResponseHeaderMaxSize = 1 << 15
ResponseStructuredMaxSize = 1 << 23
)
// the following are protocol constants
const (
responseHeaderHandlerOk = "HANDLER OK\n"
responseHeaderHandlerErrorPrefix = "HANDLER ERROR:\n"
)
+1
View File
@@ -0,0 +1 @@
package dataconn
@@ -0,0 +1,360 @@
package frameconn
import (
"encoding/binary"
"errors"
"fmt"
"io"
"net"
"sync"
"syscall"
"time"
"github.com/prometheus/client_golang/prometheus"
"github.com/zrepl/zrepl/internal/rpc/dataconn/base2bufpool"
"github.com/zrepl/zrepl/internal/rpc/dataconn/timeoutconn"
)
type FrameHeader struct {
Type uint32
PayloadLen uint32
}
// The 4 MSBs of ft are reserved for frameconn.
func IsPublicFrameType(ft uint32) bool {
return (0xf<<28)&ft == 0
}
const (
rstFrameType uint32 = 1<<28 + iota
)
func assertPublicFrameType(frameType uint32) {
if !IsPublicFrameType(frameType) {
panic(fmt.Sprintf("frameconn: frame type %v cannot be used by consumers of this package", frameType))
}
}
func (f *FrameHeader) Unmarshal(buf []byte) {
if len(buf) != 8 {
panic("frame header is 8 bytes long")
}
f.Type = binary.BigEndian.Uint32(buf[0:4])
f.PayloadLen = binary.BigEndian.Uint32(buf[4:8])
}
type Conn struct {
readMtx, writeMtx sync.Mutex
nc *timeoutconn.Conn
readNextValid bool
readNext FrameHeader
nextReadErr error
bufPool *base2bufpool.Pool // no need for sync around it
shutdown shutdownFSM
}
func Wrap(nc *timeoutconn.Conn) *Conn {
return &Conn{
nc: nc,
// ncBuf: bufio.NewReadWriter(bufio.NewReaderSize(nc, 1<<23), bufio.NewWriterSize(nc, 1<<23)),
bufPool: base2bufpool.New(15, 22, base2bufpool.Allocate), // FIXME switch to Panic, but need to enforce the limits in recv for that. => need frameconn config
readNext: FrameHeader{},
readNextValid: false,
}
}
var ErrReadFrameLengthShort = errors.New("read frame length too short")
var ErrFixedFrameLengthMismatch = errors.New("read frame length mismatch")
type Buffer struct {
bufpoolBuffer base2bufpool.Buffer
payloadLen uint32
}
func (b *Buffer) Free() {
b.bufpoolBuffer.Free()
}
func (b *Buffer) Bytes() []byte {
return b.bufpoolBuffer.Bytes()[0:b.payloadLen]
}
type Frame struct {
Header FrameHeader
Buffer Buffer
}
var ErrShutdown = fmt.Errorf("frameconn: shutting down")
// ReadFrame reads a frame from the connection.
//
// Due to an internal optimization (Readv, specifically), it is not guaranteed that a single call to
// WriteFrame unblocks a pending ReadFrame on an otherwise idle (empty) connection.
// The only way to guarantee that all previously written frames can reach the peer's layers on top
// of frameconn is to send an empty frame (no payload) and to ignore empty frames on the receiving side.
func (c *Conn) ReadFrame() (Frame, error) {
if c.shutdown.IsShuttingDown() {
return Frame{}, ErrShutdown
}
// only acquire readMtx now to prioritize the draining in Shutdown()
// over external callers (= drain public callers)
c.readMtx.Lock()
defer c.readMtx.Unlock()
f, err := c.readFrame()
if f.Header.Type == rstFrameType {
c.shutdown.Begin()
return Frame{}, ErrShutdown
}
return f, err
}
// callers must have readMtx locked
func (c *Conn) readFrame() (Frame, error) {
if c.nextReadErr != nil {
ret := c.nextReadErr
c.nextReadErr = nil
return Frame{}, ret
}
if !c.readNextValid {
var buf [8]byte
if _, err := io.ReadFull(c.nc, buf[:]); err != nil {
return Frame{}, err
}
c.readNext.Unmarshal(buf[:])
c.readNextValid = true
}
// read payload + next header
var nextHdrBuf [8]byte
buffer := c.bufPool.Get(uint(c.readNext.PayloadLen))
bufferBytes := buffer.Bytes()
if c.readNext.PayloadLen == 0 {
// This if statement implements the unlock-by-sending-empty-frame behavior
// documented in ReadFrame's public docs.
//
// It is crucial that we return this empty frame now:
// Consider the following plot with x-axis being time,
// P being a frame with payload, E one without, X either of P or E
//
// P P P P P P P E.....................X
// | | | |
// | | | F3
// | | |
// | F2 |significant time between frames because
// F1 the peer has nothing to say to us
//
// Assume we're at the point were F2's header is in c.readNext.
// That means F2 has not yet been returned.
// But because it is empty (no payload), we're already done reading it.
// If we omitted this if statement, the following would happen:
// Readv below would read [][]byte{[len(0)], [len(8)]).
c.readNextValid = false
frame := Frame{
Header: c.readNext,
Buffer: Buffer{
bufpoolBuffer: buffer,
payloadLen: c.readNext.PayloadLen, // 0
},
}
return frame, nil
}
noNextHeader := false
if n, err := c.nc.ReadvFull([][]byte{bufferBytes, nextHdrBuf[:]}); err != nil {
noNextHeader = true
zeroPayloadAndPeerClosed := n == 0 && c.readNext.PayloadLen == 0 && err == io.EOF
zeroPayloadAndNextFrameHeaderThenPeerClosed := err == io.EOF && c.readNext.PayloadLen == 0 && n == int64(len(nextHdrBuf))
nonzeroPayloadRecvdButNextHeaderMissing := n > 0 && uint32(n) == c.readNext.PayloadLen
if zeroPayloadAndPeerClosed || zeroPayloadAndNextFrameHeaderThenPeerClosed || nonzeroPayloadRecvdButNextHeaderMissing {
// This is the last frame on the conn.
// Store the error to be returned on the next invocation of ReadFrame.
c.nextReadErr = err
// NORETURN, this frame is still valid
} else {
return Frame{}, err
}
}
frame := Frame{
Header: c.readNext,
Buffer: Buffer{
bufpoolBuffer: buffer,
payloadLen: c.readNext.PayloadLen,
},
}
if !noNextHeader {
c.readNext.Unmarshal(nextHdrBuf[:])
c.readNextValid = true
} else {
c.readNextValid = false
}
return frame, nil
}
func (c *Conn) WriteFrame(payload []byte, frameType uint32) error {
assertPublicFrameType(frameType)
if c.shutdown.IsShuttingDown() {
return ErrShutdown
}
c.writeMtx.Lock()
defer c.writeMtx.Unlock()
return c.writeFrame(payload, frameType)
}
func (c *Conn) writeFrame(payload []byte, frameType uint32) error {
var hdrBuf [8]byte
binary.BigEndian.PutUint32(hdrBuf[0:4], frameType)
binary.BigEndian.PutUint32(hdrBuf[4:8], uint32(len(payload)))
bufs := net.Buffers([][]byte{hdrBuf[:], payload})
if _, err := c.nc.WritevFull(bufs); err != nil {
return err
}
return nil
}
func (c *Conn) ResetWriteTimeout() error {
return c.nc.RenewWriteDeadline()
}
func (c *Conn) Shutdown(deadline time.Time) error {
// TCP connection teardown is a bit wonky if we are in a situation
// where there is still data in flight (DIF) to our side:
// If we just close the connection, our kernel will send RSTs
// in response to the DIF, and those RSTs may reach the client's
// kernel faster than the client app is able to pull the
// last bytes from its kernel TCP receive buffer.
//
// Therefore, we send a frame with type rstFrameType to indicate
// that the connection is to be closed immediately, and further
// use CloseWrite instead of Close.
// As per definition of the wire interface, CloseWrite guarantees
// delivery of the data in our kernel TCP send buffer.
// Therefore, the client always receives the RST frame.
//
// Now what are we going to do after that?
//
// 1. Naive Option: We just call Close() right after CloseWrite.
// This yields the same race condition as explained above (DIF, first
// paragraph): The situation just became a little more unlikely because
// our rstFrameType + CloseWrite dance gave the client a full RTT worth of
// time to read the data from its TCP recv buffer.
//
// 2. Correct Option: Drain the read side until io.EOF
// We can read from the unclosed read-side of the connection until we get
// the io.EOF caused by the (well behaved) client closing the connection
// in response to it reading the rstFrameType frame we sent.
// However, this wastes resources on our side (we don't care about the
// pending data anymore), and has potential for (D)DoS through CPU-time
// exhaustion if the client just keeps sending data.
// Then again, this option has the advantage with well-behaved clients
// that we do not waste precious kernel-memory on the stale receive buffer
// on our side (which is still full of data that we do not intend to read).
//
// 2.1 DoS Mitigation: Bound the number of bytes to drain, then close
// At the time of writing, this technique is practiced by the Go http server
// implementation, and actually SHOULDed in the HTTP 1.1 RFC. It is
// important to disable the idle timeout of the underlying timeoutconn in
// that case and set an absolute deadline by which the socket must have
// been fully drained. Not too hard, though ;)
//
// 2.2: Client sends RST, not FIN when it receives an rstFrameTyp frame.
// We can use wire.(*net.TCPConn).SetLinger(0) to force an RST to be sent
// on a subsequent close (instead of a FIN + wait for FIN+ACK).
// TODO put this into Wire interface as an abstract method.
//
// 2.3 Only start draining after N*RTT
// We have an RTT approximation from Wire.CloseWrite, which by definition
// must not return before all to-be-sent-data has been acknowledged by the
// client. Give the client a fair chance to react, and only start draining
// after a multiple of the RTT has elapsed.
// We waste the recv buffer memory a little longer than necessary, iff the
// client reacts faster than expected. But we don't wast CPU time.
// If we apply 2.2, we'll also have the benefit that our kernel will have
// dropped the recv buffer memory as soon as it receives the client's RST.
//
// 3. TCP-only: OOB-messaging
// We can use TCP's 'urgent' flag in the client to acknowledge the receipt
// of the rstFrameType to us.
// We can thus wait for that signal while leaving the kernel buffer as is.
// TODO: For now, we just drain the connection (Option 2),
// but we enforce deadlines so the _time_ we drain the connection
// is bounded, although we do _that_ at full speed
defer prometheus.NewTimer(prom.ShutdownSeconds).ObserveDuration()
closeWire := func(step string) error {
// TODO SetLinger(0) or similar (we want RST frames here, not FINS)
closeErr := c.nc.Close()
if closeErr == nil {
return nil
}
// TODO go1.13: https://github.com/zrepl/zrepl/issues/190
// https://github.com/golang/go/issues/8319
// (use errors.Is(closeErr, syscall.ECONNRESET))
if pe, ok := closeErr.(*net.OpError); ok && pe.Err == syscall.ECONNRESET {
// connection reset by peer on FreeBSD, see https://github.com/zrepl/zrepl/issues/190
// We know from kernel code reading that the FD behind c.nc is closed, so let's not consider this an error
return nil
}
prom.ShutdownCloseErrors.WithLabelValues("close").Inc()
return closeErr
}
hardclose := func(err error, step string) error {
prom.ShutdownHardCloses.WithLabelValues(step).Inc()
return closeWire(step)
}
c.shutdown.Begin()
// new calls to c.ReadFrame and c.WriteFrame will now return ErrShutdown
// Acquiring writeMtx and readMtx afterwards ensures that already-running calls exit successfully
// disable renewing timeouts now, enforce the requested deadline instead
// we need to do this before acquiring locks to enforce the timeout on slow
// clients / if something hangs (DoS mitigation)
if err := c.nc.DisableTimeouts(); err != nil {
return hardclose(err, "disable_timeouts")
}
if err := c.nc.SetDeadline(deadline); err != nil {
return hardclose(err, "set_deadline")
}
c.writeMtx.Lock()
defer c.writeMtx.Unlock()
if err := c.writeFrame([]byte{}, rstFrameType); err != nil {
return hardclose(err, "write_frame")
}
if err := c.nc.CloseWrite(); err != nil {
return hardclose(err, "close_write")
}
c.readMtx.Lock()
defer c.readMtx.Unlock()
// TODO DoS mitigation: wait for client acknowledgement that they initiated Shutdown,
// then perform abortive close on our side. As explained above, probably requires
// OOB signaling such as TCP's urgent flag => transport-specific?
// TODO DoS mitigation by reading limited number of bytes
// see discussion above why this is non-trivial
defer prometheus.NewTimer(prom.ShutdownDrainSeconds).ObserveDuration()
n, _ := io.Copy(io.Discard, c.nc)
prom.ShutdownDrainBytesRead.Observe(float64(n))
return closeWire("close")
}
@@ -0,0 +1,63 @@
package frameconn
import "github.com/prometheus/client_golang/prometheus"
var prom struct {
ShutdownDrainBytesRead prometheus.Summary
ShutdownSeconds prometheus.Summary
ShutdownDrainSeconds prometheus.Summary
ShutdownHardCloses *prometheus.CounterVec
ShutdownCloseErrors *prometheus.CounterVec
}
func init() {
prom.ShutdownDrainBytesRead = prometheus.NewSummary(prometheus.SummaryOpts{
Namespace: "zrepl",
Subsystem: "frameconn",
Name: "shutdown_drain_bytes_read",
Help: "Number of bytes read during the drain phase of connection shutdown",
})
prom.ShutdownSeconds = prometheus.NewSummary(prometheus.SummaryOpts{
Namespace: "zrepl",
Subsystem: "frameconn",
Name: "shutdown_seconds",
Help: "Seconds it took for connection shutdown to complete",
})
prom.ShutdownDrainSeconds = prometheus.NewSummary(prometheus.SummaryOpts{
Namespace: "zrepl",
Subsystem: "frameconn",
Name: "shutdown_drain_seconds",
Help: "Seconds it took from read-side-drain until shutdown completion",
})
prom.ShutdownHardCloses = prometheus.NewCounterVec(prometheus.CounterOpts{
Namespace: "zrepl",
Subsystem: "frameconn",
Name: "shutdown_hard_closes",
Help: "Number of hard connection closes during shutdown (abortive close)",
}, []string{"step"})
prom.ShutdownCloseErrors = prometheus.NewCounterVec(prometheus.CounterOpts{
Namespace: "zrepl",
Subsystem: "frameconn",
Name: "shutdown_close_errors",
Help: "Number of errors closing the underlying network connection. Should alert on this",
}, []string{"step"})
}
func PrometheusRegister(registry prometheus.Registerer) error {
if err := registry.Register(prom.ShutdownDrainBytesRead); err != nil {
return err
}
if err := registry.Register(prom.ShutdownSeconds); err != nil {
return err
}
if err := registry.Register(prom.ShutdownDrainSeconds); err != nil {
return err
}
if err := registry.Register(prom.ShutdownHardCloses); err != nil {
return err
}
if err := registry.Register(prom.ShutdownCloseErrors); err != nil {
return err
}
return nil
}
@@ -0,0 +1,30 @@
package frameconn
import "sync"
type shutdownFSM struct {
mtx sync.Mutex
state shutdownFSMState
}
type shutdownFSMState uint32
const (
// zero value is important
shutdownStateOpen shutdownFSMState = iota
shutdownStateBegin
)
func (f *shutdownFSM) Begin() (thisCallStartedShutdown bool) {
f.mtx.Lock()
defer f.mtx.Unlock()
thisCallStartedShutdown = f.state != shutdownStateOpen
f.state = shutdownStateBegin
return thisCallStartedShutdown
}
func (f *shutdownFSM) IsShuttingDown() bool {
f.mtx.Lock()
defer f.mtx.Unlock()
return f.state != shutdownStateOpen
}
@@ -0,0 +1,21 @@
package frameconn
import (
"fmt"
"testing"
"github.com/stretchr/testify/assert"
)
func TestIsPublicFrameType(t *testing.T) {
for i := uint32(0); i < 256; i++ {
i := i
t.Run(fmt.Sprintf("^%d", i), func(t *testing.T) {
assert.False(t, IsPublicFrameType(^i))
})
}
assert.True(t, IsPublicFrameType(0))
assert.True(t, IsPublicFrameType(1))
assert.True(t, IsPublicFrameType(255))
assert.False(t, IsPublicFrameType(rstFrameType))
}
@@ -0,0 +1,151 @@
package heartbeatconn
import (
"fmt"
"net"
"sync/atomic"
"time"
"github.com/zrepl/zrepl/internal/rpc/dataconn/frameconn"
"github.com/zrepl/zrepl/internal/rpc/dataconn/timeoutconn"
)
type Conn struct {
state state
fc *frameconn.Conn
sendInterval, timeout time.Duration
stopSend chan struct{}
lastFrameSent atomic.Value // time.Time
}
type HeartbeatTimeout struct{}
func (e HeartbeatTimeout) Error() string {
return "heartbeat timeout"
}
// 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 HeartbeatTimeout) Temporary() bool { return true }
func (e HeartbeatTimeout) Timeout() bool { return true }
var _ net.Error = HeartbeatTimeout{}
type state = int32
const (
stateInitial state = 0
stateClosed state = 2
)
const (
heartbeat uint32 = 1 << 24
)
// The 4 MSBs of ft are reserved for frameconn, we reserve the next 4 MSB for us.
func IsPublicFrameType(ft uint32) bool {
return frameconn.IsPublicFrameType(ft) && (0xf<<24)&ft == 0
}
func assertPublicFrameType(frameType uint32) {
if !IsPublicFrameType(frameType) {
panic(fmt.Sprintf("heartbeatconn: frame type %v cannot be used by consumers of this package", frameType))
}
}
func Wrap(nc timeoutconn.Wire, sendInterval, timeout time.Duration) *Conn {
c := &Conn{
fc: frameconn.Wrap(timeoutconn.Wrap(nc, timeout)),
stopSend: make(chan struct{}),
sendInterval: sendInterval,
timeout: timeout,
}
c.lastFrameSent.Store(time.Now())
go c.sendHeartbeats()
return c
}
func (c *Conn) Shutdown() error {
normalClose := atomic.CompareAndSwapInt32(&c.state, stateInitial, stateClosed)
if normalClose {
close(c.stopSend)
}
return c.fc.Shutdown(time.Now().Add(c.timeout))
}
// started as a goroutine in constructor
func (c *Conn) sendHeartbeats() {
sleepTime := func(now time.Time) time.Duration {
lastSend := c.lastFrameSent.Load().(time.Time)
return lastSend.Add(c.sendInterval).Sub(now)
}
timer := time.NewTimer(sleepTime(time.Now()))
defer timer.Stop()
for {
select {
case <-c.stopSend:
return
case now := <-timer.C:
func() {
defer func() {
timer.Reset(sleepTime(time.Now()))
}()
if sleepTime(now) > 0 {
return
}
debug("send heartbeat")
// if the connection is in zombie mode (aka iptables DROP inbetween peers)
// this call or one of its successors will block after filling up the kernel tx buffer
err := c.fc.WriteFrame([]byte{}, heartbeat)
if err != nil {
debug("send heartbeat error: %s", err)
}
// ignore errors from WriteFrame to rate-limit SendHeartbeat retries
c.lastFrameSent.Store(time.Now())
}()
}
}
}
func (c *Conn) ReadFrame() (frameconn.Frame, error) {
return c.readFrameFiltered()
}
func (c *Conn) readFrameFiltered() (frameconn.Frame, error) {
for {
f, err := c.fc.ReadFrame()
if err != nil {
return frameconn.Frame{}, err
}
if IsPublicFrameType(f.Header.Type) {
return f, nil
}
if f.Header.Type != heartbeat {
return frameconn.Frame{}, fmt.Errorf("unknown frame type %x", f.Header.Type)
}
// drop heartbeat frame
debug("received heartbeat, resetting write timeout")
// the peer's heartbeat proves to us that the peer is still live
// => trust the peer at this point (DoS risks are ignored ATM)
// => we assume that the connection is symmetric duplex, i.e., if receiving works for us,
// sending works for us, too.
// So, let's grant the peer another write timeout.
err = c.fc.ResetWriteTimeout()
debug("renew frameconn write timeout returned errT=%T err=%s", err, err)
if err != nil {
return frameconn.Frame{}, err
}
continue
}
}
func (c *Conn) WriteFrame(payload []byte, frameType uint32) error {
assertPublicFrameType(frameType)
err := c.fc.WriteFrame(payload, frameType)
if err == nil {
c.lastFrameSent.Store(time.Now())
}
return err
}
@@ -0,0 +1,21 @@
package heartbeatconn
import (
"fmt"
"os"
)
var debugEnabled bool = false
func init() {
if os.Getenv("ZREPL_RPC_DATACONN_HEARTBEATCONN_DEBUG") != "" {
debugEnabled = true
}
}
//nolint:deadcode,unused
func debug(format string, args ...interface{}) {
if debugEnabled {
fmt.Fprintf(os.Stderr, "rpc/dataconn/heartbeatconn: %s\n", fmt.Sprintf(format, args...))
}
}
@@ -0,0 +1,27 @@
package heartbeatconn
import (
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/zrepl/zrepl/internal/rpc/dataconn/frameconn"
)
func TestFrameTypes(t *testing.T) {
assert.True(t, frameconn.IsPublicFrameType(heartbeat))
}
func TestNegativeTimer(t *testing.T) {
timer := time.NewTimer(-1 * time.Second)
defer timer.Stop()
time.Sleep(100 * time.Millisecond)
select {
case <-timer.C:
t.Log("timer with negative time fired, that's what we want")
default:
t.Fail()
}
}
@@ -0,0 +1,140 @@
// This integration test exercises the behavior of heartbeatconn
// where the server is slow at handling the data received over the connection.
// Note that the server is still sending heartbeats to the client, it's just the
// data handling (usually I/O in case of zrepl endpoint.Receiver) that is slow.
//
// In commit 082335df5d85e1b0b9faa35ff182c71886142d3e and earlier, heartbeatconn would fail
// this benchmark with a writev I/O timeout (here the ss(8) output at the time of failure)
//
// ESTAB 33369 0 127.0.0.1:12345 127.0.0.1:57282 users:(("heartbeatconn_i",pid=25953,fd=5))
// cubic wscale:7,7 rto:203 rtt:2.992/5.849 ato:162 mss:32768 pmtu:65535 rcvmss:32741 advmss:65483 cwnd:10 bytes_sent:48 bytes_acked:48 bytes_received:195401 segs_out:44 segs_in:57 data_segs_out:6 data_segs_in:34 send 876.1Mbps lastsnd:125 lastrcv:9390 lastack:125 pacing_rate 1752.0Mbps delivery_rate 6393.8Mbps delivered:7 app_limited busy:42ms rcv_rtt:1 rcv_space:65483 rcv_ssthresh:65483 minrtt:0.029
// --
// ESTAB 0 3956805 127.0.0.1:57282 127.0.0.1:12345 users:(("heartbeatconn_i",pid=26100,fd=3))
// cubic wscale:7,7 rto:211 backoff:5 rtt:10.38/16.937 ato:40 mss:32768 pmtu:65535 rcvmss:536 advmss:65483 cwnd:10 bytes_sent:195401 bytes_acked:195402 bytes_received:48 segs_out:57 segs_in:45 data_segs_out:34 data_segs_in:6 send 252.5Mbps lastsnd:9390 lastrcv:125 lastack:125 pacing_rate 505.1Mbps delivery_rate 1971.0Mbps delivered:35 busy:30127ms rwnd_limited:30086ms(99.9%) rcv_space:65495 rcv_ssthresh:65495 notsent:3956805 minrtt:0.007
// panic: writev tcp 127.0.0.1:57282->127.0.0.1:12345: i/o timeout
//
// The assumed reason for those writev timeouts is the following:
// - Sporadic server stalls (sever data handling, usually I/O) cause TCP exponential backoff on the client for client->server
// - Go runtime unblocks after the deadline expires, resultin gin writev I/O timeout
// - That is, even though the client observed heartbeats from the server
// -> TCP doesn't assume symmetric connection behavior, but our implementation does.
//
// The fix contained in the commit this message was committed with resets the deadline whenever
// a heartbeat is received from the server.
//
// How to run this integration test:
//
// Terminal 1:
// $ ZREPL_RPC_DATACONN_HEARTBEATCONN_DEBUG=1 go run heartbeatconn_integration_variablereceiverate.go -mode server -addr 127.0.0.1:12345
// rpc/dataconn/heartbeatconn: send heartbeat
// rpc/dataconn/heartbeatconn: send heartbeat
// ...
//
// Terminal 2:
// $ ZREPL_RPC_DATACONN_HEARTBEATCONN_DEBUG=1 go run heartbeatconn_integration_variablereceiverate.go -mode client -addr 127.0.0.1:12345
// rpc/dataconn/heartbeatconn: received heartbeat, resetting write timeout
// rpc/dataconn/heartbeatconn: renew frameconn write timeout returned errT=<nil> err=%!s(<nil>)
// rpc/dataconn/heartbeatconn: send heartbeat
// rpc/dataconn/heartbeatconn: received heartbeat, resetting write timeout
// rpc/dataconn/heartbeatconn: renew frameconn write timeout returned errT=<nil> err=%!s(<nil>)
// rpc/dataconn/heartbeatconn: received heartbeat, resetting write timeout
// ...
package main
import (
"bytes"
"flag"
"fmt"
"io"
"log"
"math/rand"
"net"
"os"
"os/exec"
"path"
"time"
"github.com/zrepl/zrepl/internal/util/devnoop"
"github.com/zrepl/zrepl/internal/rpc/dataconn/heartbeatconn"
)
func orDie(err error) {
if err != nil {
grepfield := path.Base(os.Args[0])[:10]
fmt.Fprintf(os.Stderr, "grepping for %s\n", grepfield)
sh := fmt.Sprintf("ss -ntpi | grep -A1 %s", grepfield)
cmd := exec.Command("bash", "-c", sh)
o, _ := cmd.CombinedOutput()
buf := bytes.NewBuffer(o)
_, _ = io.Copy(os.Stderr, buf)
panic(err)
}
}
var mode string
var addr string
func main() {
flag.StringVar(&mode, "mode", "", "server|client")
flag.StringVar(&addr, "addr", "INVALID", "")
flag.Parse()
modemap := map[string]func(){
"server": server,
"client": client,
}
modemap[mode]()
}
func server() {
ln, err := net.Listen("tcp", addr)
orDie(err)
l := ln.(*net.TCPListener)
for {
c, err := l.AcceptTCP()
if err != nil {
log.Printf("accept err: %s", err)
continue
}
hc := heartbeatconn.Wrap(c, 5*time.Second, 10*time.Second)
for {
f, err := hc.ReadFrame()
orDie(err)
// _, err = buf.Write(f.Buffer.Bytes())
// orDie(err)
sleep := time.Duration(rand.NormFloat64()*500) * time.Millisecond
time.Sleep(sleep)
f.Buffer.Free()
}
}
}
func client() {
c, err := net.Dial("tcp", addr)
orDie(err)
hc := heartbeatconn.Wrap(c.(*net.TCPConn), 5*time.Second, 10*time.Second)
// follow API requirements to always ReadFrame
go func() {
for {
f, err := hc.ReadFrame()
orDie(err)
f.Buffer.Free()
}
}()
dn := devnoop.Get()
var buf [1 << 10]byte
for {
n, err := dn.Read(buf[:])
orDie(err)
err = hc.WriteFrame(buf[:n], 23)
orDie(err)
}
}
@@ -0,0 +1,173 @@
// microbenchmark to manually test rpc/dataconn performance
//
// With stdin / stdout on client and server, simulating zfs send|recv piping
//
// ./microbenchmark -appmode server | pv -r > /dev/null
// ./microbenchmark -appmode client -direction recv < /dev/zero
//
// Without the overhead of pipes (just protocol performance, mostly useful with perf bc no bw measurement)
//
// ./microbenchmark -appmode client -direction recv -devnoopWriter -devnoopReader
// ./microbenchmark -appmode server -devnoopReader -devnoopWriter
package main
import (
"context"
"flag"
"fmt"
"io"
"net"
"os"
"github.com/pkg/profile"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/replication/logic/pdu"
"github.com/zrepl/zrepl/internal/rpc/dataconn"
"github.com/zrepl/zrepl/internal/rpc/dataconn/timeoutconn"
"github.com/zrepl/zrepl/internal/transport"
"github.com/zrepl/zrepl/internal/util/devnoop"
)
func orDie(err error) {
if err != nil {
panic(err)
}
}
type readerStreamCopier struct{ io.Reader }
func (readerStreamCopier) Close() error { return nil }
type devNullHandler struct{}
func (devNullHandler) Send(ctx context.Context, r *pdu.SendReq) (*pdu.SendRes, io.ReadCloser, error) {
var res pdu.SendRes
if args.devnoopReader {
return &res, readerStreamCopier{devnoop.Get()}, nil
} else {
return &res, readerStreamCopier{os.Stdin}, nil
}
}
func (devNullHandler) Receive(ctx context.Context, r *pdu.ReceiveReq, stream io.ReadCloser) (*pdu.ReceiveRes, error) {
var out io.Writer = os.Stdout
if args.devnoopWriter {
out = devnoop.Get()
}
_, err := io.Copy(out, stream)
var res pdu.ReceiveRes
return &res, err
}
func (devNullHandler) PingDataconn(ctx context.Context, r *pdu.PingReq) (*pdu.PingRes, error) {
return &pdu.PingRes{
Echo: r.GetMessage(),
}, nil
}
type tcpConnecter struct {
addr string
}
func (c tcpConnecter) Connect(ctx context.Context) (timeoutconn.Wire, error) {
conn, err := net.Dial("tcp", c.addr)
if err != nil {
return nil, err
}
return conn.(*net.TCPConn), nil
}
type tcpListener struct {
nl *net.TCPListener
clientIdent string
}
func (l tcpListener) Accept(ctx context.Context) (*transport.AuthConn, error) {
tcpconn, err := l.nl.AcceptTCP()
orDie(err)
return transport.NewAuthConn(tcpconn, l.clientIdent), nil
}
func (l tcpListener) Addr() net.Addr { return l.nl.Addr() }
func (l tcpListener) Close() error { return l.nl.Close() }
var args struct {
addr string
appmode string
direction string
profile bool
devnoopReader bool
devnoopWriter bool
}
func server() {
log := logger.NewStderrDebugLogger()
log.Debug("starting server")
nl, err := net.Listen("tcp", args.addr)
orDie(err)
l := tcpListener{nl.(*net.TCPListener), "fakeclientidentity"}
srv := dataconn.NewServer(nil, nil, logger.NewStderrDebugLogger(), devNullHandler{})
ctx := context.Background()
srv.Serve(ctx, l)
}
func main() {
flag.BoolVar(&args.profile, "profile", false, "")
flag.BoolVar(&args.devnoopReader, "devnoopReader", false, "")
flag.BoolVar(&args.devnoopWriter, "devnoopWriter", false, "")
flag.StringVar(&args.addr, "address", ":8888", "")
flag.StringVar(&args.appmode, "appmode", "client|server", "")
flag.StringVar(&args.direction, "direction", "", "send|recv")
flag.Parse()
if args.profile {
defer profile.Start(profile.CPUProfile).Stop()
}
switch args.appmode {
case "client":
client()
case "server":
server()
default:
orDie(fmt.Errorf("unknown appmode %q", args.appmode))
}
}
func client() {
logger := logger.NewStderrDebugLogger()
ctx := context.Background()
connecter := tcpConnecter{args.addr}
client := dataconn.NewClient(connecter, logger)
switch args.direction {
case "send":
req := pdu.SendReq{}
_, stream, err := client.ReqSend(ctx, &req)
orDie(err)
_, err = io.Copy(os.Stdout, stream)
orDie(err)
case "recv":
var r io.Reader = os.Stdin
if args.devnoopReader {
r = devnoop.Get()
}
s := readerStreamCopier{r}
req := pdu.ReceiveReq{}
_, err := client.ReqRecv(ctx, &req, &s)
orDie(err)
default:
orDie(fmt.Errorf("unknown direction%q", args.direction))
}
}
+289
View File
@@ -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)}
}
+296
View File
@@ -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...))
}
}
+132
View File
@@ -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()
}
@@ -0,0 +1,12 @@
# setup-specific
inventory
*.retry
# generated by gen_files.sh
files/*ssh_client_identity
files/*ssh_client_identity.pub
files/*.tls.*.key
files/*.tls.*.csr
files/*.tls.*.crt
files/wireevaluator
@@ -0,0 +1,15 @@
This directory contains very hacky test automation for wireevaluator based on nested Ansible playbooks.
* Copy `inventory.example` to `inventory`
* Adjust `inventory` IP addresses as needed
* Make sure there's an OpenSSH server running on the serve host
* Make sure there's no firewalling whatsoever between the hosts
* Run `GENKEYS=1 ./gen_files.sh` to re-generate self-signed TLS certs
* Run the following command, adjusting the `wireevaluator_repeat` value to the number of times you want to repeat each test
```
ansible-playbook -i inventory all.yml -e `wireevaluator_repeat=3`
```
Generally, things are fine if the playbook doesn't show any panics from wireevaluator.
@@ -0,0 +1,17 @@
- hosts: connect,serve
tasks:
- name: "run test"
include: internal_prepare_and_run_repeated.yml
wireevaluator_transport: "{{config.0}}"
wireevaluator_case: "{{config.1}}"
wireevaluator_repeat: "{{wireevaluator_repeat}}"
with_cartesian:
- [ tls, ssh, tcp ]
-
- closewrite_server
- closewrite_client
- readdeadline_server
- readdeadline_client
loop_control:
loop_var: config
@@ -0,0 +1,55 @@
#!/usr/bin/env bash
set -e
cd "$( dirname "${BASH_SOURCE[0]}")"
FILESDIR="$(pwd)"/files
echo "[INFO] compile binary"
pushd .. >/dev/null
go build -o $FILESDIR/wireevaluator
popd >/dev/null
if [ "$GENKEYS" == "" ]; then
echo "[INFO] GENKEYS environment variable not set, assumed to be valid"
exit 0
fi
echo "[INFO] gen ssh key"
ssh-keygen -f "$FILESDIR/wireevaluator.ssh_client_identity" -t ed25519
echo "[INFO] gen tls keys"
cakey="$FILESDIR/wireevaluator.tls.ca.key"
cacrt="$FILESDIR/wireevaluator.tls.ca.crt"
hostprefix="$FILESDIR/wireevaluator.tls"
openssl genrsa -out "$cakey" 4096
openssl req -x509 -new -nodes -key "$cakey" -sha256 -days 1 -out "$cacrt"
declare -a HOSTS
HOSTS+=("theserver")
HOSTS+=("theclient")
for host in "${HOSTS[@]}"; do
key="${hostprefix}.${host}.key"
csr="${hostprefix}.${host}.csr"
crt="${hostprefix}.${host}.crt"
openssl genrsa -out "$key" 2048
(
echo "."
echo "."
echo "."
echo "."
echo "."
echo $host
echo "."
echo "."
echo "."
echo "."
) | openssl req -new -key "$key" -out "$csr"
openssl x509 -req -in "$csr" -CA "$cacrt" -CAkey "$cakey" -CAcreateserial -out "$crt" -days 1 -sha256
done
@@ -0,0 +1,54 @@
---
- name: compile binary and any key files required
local_action: command ./gen_files.sh
- name: Kill test binary
shell: "killall -9 wireevaluator || true"
- name: Deploy new binary
copy:
src: "files/wireevaluator"
dest: "/opt/wireevaluator"
mode: 0755
- set_fact:
wireevaluator_connect_ip: "{{hostvars['connect'].ansible_host}}"
wireevaluator_serve_ip: "{{hostvars['serve'].ansible_host}}"
- name: Deploy config
template:
src: "templates/{{wireevaluator_transport}}.yml.j2"
dest: "/opt/wireevaluator.yml"
- name: Deploy client identity
copy:
src: "files/wireevaluator.{{item}}"
dest: "/opt/wireevaluator.{{item}}"
mode: 0400
with_items:
- ssh_client_identity
- ssh_client_identity.pub
- tls.ca.key
- tls.ca.crt
- tls.theserver.key
- tls.theserver.crt
- tls.theclient.key
- tls.theclient.crt
- name: Setup server ssh client identity access
when: inventory_hostname == "serve"
block:
- authorized_key:
user: root
state: present
key: "{{ lookup('file', 'files/wireevaluator.ssh_client_identity.pub') }}"
key_options: 'command="/opt/wireevaluator -mode stdinserver -config /opt/wireevaluator.yml client1"'
- file:
state: directory
mode: 0700
path: /tmp/wireevaluator_stdinserver
- name: repeated test
include: internal_run_test_prepared_single.yml
with_sequence: start=1 end={{wireevaluator_repeat}}
@@ -0,0 +1,38 @@
---
- debug:
msg: "run test transport={{wireevaluator_transport}} case={{wireevaluator_case}} repeatedly"
- name: Run Server
when: inventory_hostname == "serve"
command: /opt/wireevaluator -config /opt/wireevaluator.yml -mode serve -testcase {{wireevaluator_case}}
register: spawn_servers
async: 60
poll: 0
- name: Run Client
when: inventory_hostname == "connect"
command: /opt/wireevaluator -config /opt/wireevaluator.yml -mode connect -testcase {{wireevaluator_case}}
register: spawn_clients
async: 60
poll: 0
- name: Wait for server shutdown
when: inventory_hostname == "serve"
async_status:
jid: "{{ spawn_servers.ansible_job_id}}"
delay: 0.5
retries: 10
- name: Wait for client shutdown
when: inventory_hostname == "connect"
async_status:
jid: "{{ spawn_clients.ansible_job_id}}"
delay: 0.5
retries: 10
- name: Wait for connections to die (TIME_WAIT conns)
command: sleep 4
changed_when: false
@@ -0,0 +1,2 @@
connect ansible_user=root ansible_host=192.168.122.128 wireevaluator_mode="connect"
serve ansible_user=root ansible_host=192.168.122.129 wireevaluator_mode="serve"
@@ -0,0 +1,13 @@
connect:
type: ssh+stdinserver
host: {{wireevaluator_serve_ip}}
user: root
port: 22
identity_file: /opt/wireevaluator.ssh_client_identity
options: # optional, default [], `-o` arguments passed to ssh
- "Compression=yes"
serve:
type: stdinserver
client_identities:
- "client1"
@@ -0,0 +1,10 @@
connect:
type: tcp
address: "{{wireevaluator_serve_ip}}:8888"
serve:
type: tcp
listen: ":8888"
clients: {
"{{wireevaluator_connect_ip}}" : "client1"
}
@@ -0,0 +1,16 @@
connect:
type: tls
address: "{{wireevaluator_serve_ip}}:8888"
ca: "/opt/wireevaluator.tls.ca.crt"
cert: "/opt/wireevaluator.tls.theclient.crt"
key: "/opt/wireevaluator.tls.theclient.key"
server_cn: "theserver"
serve:
type: tls
listen: ":8888"
ca: "/opt/wireevaluator.tls.ca.crt"
cert: "/opt/wireevaluator.tls.theserver.crt"
key: "/opt/wireevaluator.tls.theserver.key"
client_cns:
- "theclient"
@@ -0,0 +1,111 @@
// a tool to test whether a given transport implements the timeoutconn.Wire interface
package main
import (
"context"
"flag"
"fmt"
"os"
"path"
netssh "github.com/problame/go-netssh"
"github.com/zrepl/yaml-config"
"github.com/zrepl/zrepl/internal/config"
"github.com/zrepl/zrepl/internal/transport"
transportconfig "github.com/zrepl/zrepl/internal/transport/fromconfig"
)
func noerror(err error) {
if err != nil {
panic(err)
}
}
type Config struct {
Connect config.ConnectEnum
Serve config.ServeEnum
}
var args struct {
mode string
configPath string
testCase string
}
var conf Config
type TestCase interface {
Client(wire transport.Wire)
Server(wire transport.Wire)
}
func main() {
flag.StringVar(&args.mode, "mode", "", "connect|serve")
flag.StringVar(&args.configPath, "config", "", "config file path")
flag.StringVar(&args.testCase, "testcase", "", "")
flag.Parse()
bytes, err := os.ReadFile(args.configPath)
noerror(err)
err = yaml.UnmarshalStrict(bytes, &conf)
noerror(err)
ctx, cancel := context.WithCancel(context.Background())
defer cancel()
global := &config.Global{
Serve: &config.GlobalServe{
StdinServer: &config.GlobalStdinServer{
SockDir: "/tmp/wireevaluator_stdinserver",
},
},
}
switch args.mode {
case "connect":
tc, err := getTestCase(args.testCase)
noerror(err)
connecter, err := transportconfig.ConnecterFromConfig(global, conf.Connect, config.ParseFlagsNone)
noerror(err)
wire, err := connecter.Connect(ctx)
noerror(err)
tc.Client(wire)
case "serve":
tc, err := getTestCase(args.testCase)
noerror(err)
lf, err := transportconfig.ListenerFactoryFromConfig(global, conf.Serve, config.ParseFlagsNone)
noerror(err)
l, err := lf()
noerror(err)
conn, err := l.Accept(ctx)
noerror(err)
tc.Server(conn)
case "stdinserver":
identity := flag.Arg(0)
unixaddr := path.Join(global.Serve.StdinServer.SockDir, identity)
err := netssh.Proxy(ctx, unixaddr)
if err == nil {
os.Exit(0)
}
panic(err)
default:
panic(fmt.Sprintf("unknown mode %q", args.mode))
}
}
func getTestCase(tcName string) (TestCase, error) {
switch tcName {
case "closewrite_server":
return &CloseWrite{mode: CloseWriteServerSide}, nil
case "closewrite_client":
return &CloseWrite{mode: CloseWriteClientSide}, nil
case "readdeadline_client":
return &Deadlines{mode: DeadlineModeClientTimeout}, nil
case "readdeadline_server":
return &Deadlines{mode: DeadlineModeServerTimeout}, nil
default:
return nil, fmt.Errorf("unknown test case %q", tcName)
}
}
@@ -0,0 +1,106 @@
package main
import (
"bytes"
"io"
"log"
"github.com/zrepl/zrepl/internal/transport"
)
type CloseWriteMode uint
const (
CloseWriteClientSide CloseWriteMode = 1 + iota
CloseWriteServerSide
)
type CloseWrite struct {
mode CloseWriteMode
}
// sent repeatedly
var closeWriteTestSendData = bytes.Repeat([]byte{0x23, 0x42}, 1<<24)
var closeWriteErrorMsg = []byte{0xb, 0xa, 0xd, 0xf, 0x0, 0x0, 0xd}
func (m CloseWrite) Client(wire transport.Wire) {
switch m.mode {
case CloseWriteClientSide:
m.receiver(wire)
case CloseWriteServerSide:
m.sender(wire)
default:
panic(m.mode)
}
}
func (m CloseWrite) Server(wire transport.Wire) {
switch m.mode {
case CloseWriteClientSide:
m.sender(wire)
case CloseWriteServerSide:
m.receiver(wire)
default:
panic(m.mode)
}
}
func (CloseWrite) sender(wire transport.Wire) {
defer func() {
closeErr := wire.Close()
log.Printf("closeErr=%T %s", closeErr, closeErr)
}()
writeDone := make(chan struct{}, 1)
go func() {
close(writeDone)
for {
_, err := wire.Write(closeWriteTestSendData)
if err != nil {
return
}
}
}()
defer func() {
<-writeDone
}()
var respBuf bytes.Buffer
_, err := io.Copy(&respBuf, wire)
if err != nil {
log.Fatalf("should have received io.EOF, which is masked by io.Copy, got: %s", err)
}
if !bytes.Equal(respBuf.Bytes(), closeWriteErrorMsg) {
log.Fatalf("did not receive error message, got response with len %v:\n%v", respBuf.Len(), respBuf.Bytes())
}
}
func (CloseWrite) receiver(wire transport.Wire) {
// consume half the test data, then detect an error, send it and CloseWrite
r := io.LimitReader(wire, int64(5*len(closeWriteTestSendData)/3))
_, err := io.Copy(io.Discard, r)
noerror(err)
var errBuf bytes.Buffer
errBuf.Write(closeWriteErrorMsg)
_, err = io.Copy(wire, &errBuf)
noerror(err)
err = wire.CloseWrite()
noerror(err)
// drain wire, as documented in transport.Wire, this is the only way we know the client closed the conn
_, err = io.Copy(io.Discard, wire)
if err != nil {
// io.Copy masks io.EOF to nil, and we expect io.EOF from the client's Close() call
log.Panicf("unexpected error returned from reading conn: %s", err)
}
closeErr := wire.Close()
log.Printf("closeErr=%T %s", closeErr, closeErr)
}
@@ -0,0 +1,138 @@
package main
import (
"bytes"
"fmt"
"io"
"log"
"net"
"time"
"github.com/zrepl/zrepl/internal/transport"
)
type DeadlineMode uint
const (
DeadlineModeClientTimeout DeadlineMode = 1 + iota
DeadlineModeServerTimeout
)
type Deadlines struct {
mode DeadlineMode
}
func (d Deadlines) Client(wire transport.Wire) {
switch d.mode {
case DeadlineModeClientTimeout:
d.sleepThenSend(wire)
case DeadlineModeServerTimeout:
d.sendThenRead(wire)
default:
panic(d.mode)
}
}
func (d Deadlines) Server(wire transport.Wire) {
switch d.mode {
case DeadlineModeClientTimeout:
d.sendThenRead(wire)
case DeadlineModeServerTimeout:
d.sleepThenSend(wire)
default:
panic(d.mode)
}
}
var deadlinesTimeout = 1 * time.Second
func (d Deadlines) sleepThenSend(wire transport.Wire) {
defer wire.Close()
log.Print("sleepThenSend")
// exceed timeout of peer (do not respond to their hi msg)
time.Sleep(3 * deadlinesTimeout)
// expect that the client has hung up on us by now
err := d.sendMsg(wire, "hi")
log.Printf("err=%s", err)
log.Printf("err=%#v", err)
if err == nil {
log.Panic("no error")
}
if _, ok := err.(net.Error); !ok {
log.Panic("not a net error")
}
}
func (d Deadlines) sendThenRead(wire transport.Wire) {
log.Print("sendThenRead")
err := d.sendMsg(wire, "hi")
noerror(err)
err = wire.SetReadDeadline(time.Now().Add(deadlinesTimeout))
noerror(err)
m, err := d.recvMsg(wire)
log.Printf("m=%q", m)
log.Printf("err=%s", err)
log.Printf("err=%#v", err)
// close asap so that the peer get's a 'connection reset by peer' error or similar
closeErr := wire.Close()
if closeErr != nil {
panic(closeErr)
}
var neterr net.Error
var ok bool
if err == nil {
goto unexpErr // works for nil, too
}
neterr, ok = err.(net.Error)
if !ok {
log.Println("not a net error")
goto unexpErr
}
if !neterr.Timeout() {
log.Println("not a timeout")
}
return
unexpErr:
panic(fmt.Sprintf("sendThenRead: client should have hung up but got error %T %s", err, err))
}
const deadlinesMsgLen = 40
func (d Deadlines) sendMsg(wire transport.Wire, msg string) error {
if len(msg) > deadlinesMsgLen {
panic(len(msg))
}
var buf [deadlinesMsgLen]byte
copy(buf[:], []byte(msg))
n, err := wire.Write(buf[:])
if err != nil {
return err
}
if n != len(buf) {
panic("short write not allowed")
}
return nil
}
func (d Deadlines) recvMsg(wire transport.Wire) (string, error) {
var buf bytes.Buffer
r := io.LimitReader(wire, deadlinesMsgLen)
_, err := io.Copy(&buf, r)
if err != nil {
return "", err
}
return buf.String(), nil
}
@@ -0,0 +1,210 @@
// package timeoutconn wraps a Wire to provide idle timeouts
// based on Set{Read,Write}Deadline.
// Additionally, it exports abstractions for vectored I/O.
package timeoutconn
// NOTE
// Readv and Writev are not split-off into a separate package
// because we use raw syscalls, bypassing Conn's Read / Write methods.
import (
"errors"
"io"
"net"
"sync"
"syscall"
"time"
)
type Wire interface {
net.Conn
// A call to CloseWrite indicates that no further Write calls will be made to Wire.
// The implementation must return an error in case of Write calls after CloseWrite.
// On the peer's side, after it read all data written to Wire prior to the call to
// CloseWrite on our side, the peer's Read calls must return io.EOF.
// CloseWrite must not affect the read-direction of Wire: specifically, the
// peer must continue to be able to send, and our side must continue be
// able to receive data over Wire.
//
// Note that CloseWrite may (and most likely will) return sooner than the
// peer having received all data written to Wire prior to CloseWrite.
// Note further that buffering happening in the network stacks on either side
// mandates an explicit acknowledgement from the peer that the connection may
// be fully shut down: If we call Close without such acknowledgement, any data
// from peer to us that was already in flight may cause connection resets to
// be sent from us to the peer via the specific transport protocol. Those
// resets (e.g. RST frames) may erase all connection context on the peer,
// including data in its receive buffers. Thus, those resets are in race with
// a) transmission of data written prior to CloseWrite and
// b) the peer application reading from those buffers.
//
// The WaitForPeerClose method can be used to wait for connection termination,
// iff the implementation supports it. If it does not, the only reliable way
// to wait for a peer to have read all data from Wire (until io.EOF), is to
// expect it to close the wire at that point as well, and to drain Wire until
// we also read io.EOF.
CloseWrite() error
// Wait for the peer to close the connection.
// No data that could otherwise be Read is lost as a consequence of this call.
// The use case for this API is abortive connection shutdown.
// To provide any value over draining Wire using io.Read, an implementation
// will likely use out-of-band messaging mechanisms.
// TODO WaitForPeerClose() (supported bool, err error)
}
type Conn struct {
// immutable state
Wire
idleTimeout time.Duration
// mutable state (protected by mtx)
mtx sync.RWMutex
renewDeadlinesDisabled bool
}
func Wrap(conn Wire, idleTimeout time.Duration) *Conn {
return &Conn{
Wire: conn,
idleTimeout: idleTimeout,
}
}
// DisableTimeouts disables the idle timeout behavior provided by this package.
// Existing deadlines are cleared iff the call is the first call to this method
// or if the previous call produced an error.
func (c *Conn) DisableTimeouts() error {
c.mtx.Lock()
defer c.mtx.Unlock()
if c.renewDeadlinesDisabled {
return nil
}
err := c.SetDeadline(time.Time{})
if err != nil {
return err
}
c.renewDeadlinesDisabled = true
return nil
}
func (c *Conn) renewReadDeadline() error {
c.mtx.RLock()
defer c.mtx.RUnlock()
if c.renewDeadlinesDisabled {
return nil
}
return c.SetReadDeadline(time.Now().Add(c.idleTimeout))
}
func (c *Conn) RenewWriteDeadline() error {
c.mtx.RLock()
defer c.mtx.RUnlock()
if c.renewDeadlinesDisabled {
return nil
}
return c.SetWriteDeadline(time.Now().Add(c.idleTimeout))
}
func (c *Conn) Read(p []byte) (n int, _ error) {
n = 0
restart:
if err := c.renewReadDeadline(); err != nil {
return n, err
}
var nCurRead int
nCurRead, err := c.Wire.Read(p[n:])
n += nCurRead
if netErr, ok := err.(net.Error); ok && netErr.Timeout() && nCurRead > 0 {
goto restart
}
return n, err
}
func (c *Conn) Write(p []byte) (n int, _ error) {
n = 0
restart:
if err := c.RenewWriteDeadline(); err != nil {
return n, err
}
var nCurWrite int
nCurWrite, err := c.Wire.Write(p[n:])
n += nCurWrite
if netErr, ok := err.(net.Error); ok && netErr.Timeout() && nCurWrite > 0 {
goto restart
}
return n, err
}
// Writes the given buffers to Conn, following the semantics of io.Copy,
// but is guaranteed to use the writev system call if the wrapped Wire
// support it.
// Note the Conn does not support writev through io.Copy(aConn, aNetBuffers).
func (c *Conn) WritevFull(bufs net.Buffers) (n int64, _ error) {
n = 0
restart:
if err := c.RenewWriteDeadline(); err != nil {
return n, err
}
var nCurWrite int64
nCurWrite, err := io.Copy(c.Wire, &bufs)
n += nCurWrite
if netErr, ok := err.(net.Error); ok && netErr.Timeout() && nCurWrite > 0 {
goto restart
}
return n, err
}
var SyscallConnNotSupported = errors.New("SyscallConn not supported")
// The interface that must be implemented for vectored I/O support.
// If the wrapped Wire does not implement it, a less efficient
// fallback implementation is used.
// Rest assured that Go's *net.TCPConn implements this interface.
type SyscallConner interface {
// The sentinel error value SyscallConnNotSupported can be returned
// if the support for SyscallConn depends on runtime conditions and
// that runtime condition is not met.
SyscallConn() (syscall.RawConn, error)
}
var _ SyscallConner = (*net.TCPConn)(nil)
// Reads the given buffers full:
// Think of io.ReadvFull, but for net.Buffers + using the readv syscall.
//
// If the underlying Wire is not a SyscallConner, a fallback
// implementation based on repeated Conn.Read invocations is used.
//
// If the connection returned io.EOF, the number of bytes written until
// then + io.EOF is returned. This behavior is different to io.ReadFull
// which returns io.ErrUnexpectedEOF.
func (c *Conn) ReadvFull(buffers net.Buffers) (n int64, err error) {
return c.readv(buffers)
}
// invoked by c.readv if readv system call cannot be used
func (c *Conn) readvFallback(nbuffers net.Buffers) (n int64, err error) {
buffers := [][]byte(nbuffers)
for i := range buffers {
curBuf := buffers[i]
inner:
for len(curBuf) > 0 {
if err := c.renewReadDeadline(); err != nil {
return n, err
}
var oneN int
oneN, err = c.Read(curBuf[:]) // WE WANT NO SHADOWING
curBuf = curBuf[oneN:]
n += int64(oneN)
if err != nil {
if netErr, ok := err.(net.Error); ok && netErr.Timeout() && oneN > 0 {
continue inner
}
return n, err
}
}
}
return n, nil
}
@@ -0,0 +1,12 @@
//go:build illumos || solaris
// +build illumos solaris
package timeoutconn
import "net"
func (c *Conn) readv(buffers net.Buffers) (n int64, err error) {
// Go does not expose the SYS_READV symbol for Solaris / Illumos - do they have it?
// Anyhow, use the fallback
return c.readvFallback(buffers)
}
@@ -0,0 +1,195 @@
package timeoutconn
import (
"bytes"
"io"
"net"
"sync"
"syscall"
"testing"
"time"
"unsafe"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/zrepl/zrepl/internal/util/socketpair"
"github.com/zrepl/zrepl/internal/util/zreplcircleci"
)
func TestReadTimeout(t *testing.T) {
a, b, err := socketpair.SocketPair()
require.NoError(t, err)
defer a.Close()
defer b.Close()
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
var buf bytes.Buffer
buf.WriteString("tooktoolong")
time.Sleep(500 * time.Millisecond)
_, err := io.Copy(a, &buf)
require.NoError(t, err)
}()
go func() {
defer wg.Done()
conn := Wrap(b, 100*time.Millisecond)
buf := [4]byte{} // shorter than message put on wire
n, err := conn.Read(buf[:])
assert.Equal(t, 0, n)
assert.Error(t, err)
netErr, ok := err.(net.Error)
require.True(t, ok)
assert.True(t, netErr.Timeout())
}()
wg.Wait()
}
type writeBlockConn struct {
net.Conn
blockTime time.Duration
}
func (c writeBlockConn) Write(p []byte) (int, error) {
time.Sleep(c.blockTime)
return c.Conn.Write(p)
}
func (c writeBlockConn) CloseWrite() error {
return c.Conn.Close()
}
func TestWriteTimeout(t *testing.T) {
a, b, err := socketpair.SocketPair()
require.NoError(t, err)
defer a.Close()
defer b.Close()
var buf bytes.Buffer
buf.WriteString("message")
blockConn := writeBlockConn{a, 500 * time.Millisecond}
conn := Wrap(blockConn, 100*time.Millisecond)
n, err := conn.Write(buf.Bytes())
assert.Equal(t, 0, n)
assert.Error(t, err)
netErr, ok := err.(net.Error)
require.True(t, ok)
assert.True(t, netErr.Timeout())
}
func TestNoPartialReadsDueToDeadline(t *testing.T) {
zreplcircleci.SkipOnCircleCI(t, "needs predictable low scheduling latency")
a, b, err := socketpair.SocketPair()
require.NoError(t, err)
defer a.Close()
defer b.Close()
var wg sync.WaitGroup
wg.Add(2)
go func() {
defer wg.Done()
a.Write([]byte{1, 2, 3, 4, 5})
// sleep to provoke a partial read in the consumer goroutine
time.Sleep(50 * time.Millisecond)
a.Write([]byte{6, 7, 8, 9, 10})
}()
go func() {
defer wg.Done()
bc := Wrap(b, 100*time.Millisecond)
var buf bytes.Buffer
beginRead := time.Now()
// io.Copy will encounter a partial read, then wait ~50ms until the other 5 bytes are written
// It is still going to fail with deadline err because it expects EOF
n, err := io.Copy(&buf, bc)
readDuration := time.Since(beginRead)
t.Logf("read duration=%s", readDuration)
t.Logf("recv done n=%v err=%v", n, err)
t.Logf("buf=%v", buf.Bytes())
neterr, ok := err.(net.Error)
require.True(t, ok)
assert.True(t, neterr.Timeout())
assert.Equal(t, int64(10), n)
assert.Equal(t, []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}, buf.Bytes())
// 50ms for the second read, 100ms after that one for the deadline
// allow for some jitter
assert.True(t, readDuration > 140*time.Millisecond)
assert.True(t, readDuration < 200*time.Millisecond)
}()
wg.Wait()
}
type partialWriteMockConn struct {
net.Conn // to satisfy interface
buf bytes.Buffer
writeDuration time.Duration
returnAfterBytesWritten int
}
func newPartialWriteMockConn(writeDuration time.Duration, returnAfterBytesWritten int) *partialWriteMockConn {
return &partialWriteMockConn{
writeDuration: writeDuration,
returnAfterBytesWritten: returnAfterBytesWritten,
}
}
func (c *partialWriteMockConn) Write(p []byte) (int, error) {
time.Sleep(c.writeDuration)
consumeBytes := len(p)
if consumeBytes > c.returnAfterBytesWritten {
consumeBytes = c.returnAfterBytesWritten
}
n, err := c.buf.Write(p[0:consumeBytes])
if err != nil || n != consumeBytes {
panic("bytes.Buffer behaves unexpectedly")
}
return n, nil
}
func TestPartialWriteMockConn(t *testing.T) {
zreplcircleci.SkipOnCircleCI(t, "because it relies on scheduler responsiveness < 50ms")
mc := newPartialWriteMockConn(100*time.Millisecond, 5)
buf := []byte{1, 2, 3, 4, 5, 6, 7, 8, 9, 10}
begin := time.Now()
n, err := mc.Write(buf[:])
duration := time.Since(begin)
assert.NoError(t, err)
assert.Equal(t, 5, n)
assert.True(t, duration > 100*time.Millisecond)
assert.True(t, duration < 150*time.Millisecond)
}
func TestNoPartialWritesDueToDeadline(t *testing.T) {
a, b, err := socketpair.SocketPair()
require.NoError(t, err)
defer a.Close()
defer b.Close()
var buf bytes.Buffer
buf.WriteString("message")
blockConn := writeBlockConn{a, 150 * time.Millisecond}
conn := Wrap(blockConn, 100*time.Millisecond)
n, err := conn.Write(buf.Bytes())
assert.Equal(t, 0, n)
assert.Error(t, err)
netErr, ok := err.(net.Error)
require.True(t, ok)
assert.True(t, netErr.Timeout())
}
func TestIovecLenFieldIsMachineUint(t *testing.T) {
iov := syscall.Iovec{}
_ = iov // make linter happy (unsafe.Sizeof not recognized as usage)
size_t := unsafe.Sizeof(iov.Len)
if size_t != unsafe.Sizeof(uint(23)) {
t.Fatalf("expecting (struct iov)->Len to be sizeof(uint)")
}
// ssize_t is defined to be the signed version of size_t,
// so we know sizeof(ssize_t) == sizeof(int)
}
@@ -0,0 +1,132 @@
//go:build !illumos && !solaris
// +build !illumos,!solaris
package timeoutconn
import (
"io"
"net"
"syscall"
"unsafe"
)
func buildIovecs(buffers net.Buffers) (totalLen int64, vecs []syscall.Iovec) {
vecs = make([]syscall.Iovec, 0, len(buffers))
for i := range buffers {
totalLen += int64(len(buffers[i]))
if len(buffers[i]) == 0 {
continue
}
v := syscall.Iovec{
Base: &buffers[i][0],
}
// syscall.Iovec.Len has platform-dependent size, thus use SetLen
v.SetLen(len(buffers[i]))
vecs = append(vecs, v)
}
return totalLen, vecs
}
func (c *Conn) readv(buffers net.Buffers) (n int64, err error) {
scc, ok := c.Wire.(SyscallConner)
if !ok {
return c.readvFallback(buffers)
}
rawConn, err := scc.SyscallConn()
if err == SyscallConnNotSupported {
return c.readvFallback(buffers)
}
if err != nil {
return 0, err
}
_, iovecs := buildIovecs(buffers)
for len(iovecs) > 0 {
if err := c.renewReadDeadline(); err != nil {
return n, err
}
oneN, oneErr := c.doOneReadv(rawConn, &iovecs)
n += oneN
if netErr, ok := oneErr.(net.Error); ok && netErr.Timeout() && oneN > 0 { // TODO likely not working
continue
} else if oneErr == nil && oneN > 0 {
continue
} else {
return n, oneErr
}
}
return n, nil
}
func (c *Conn) doOneReadv(rawConn syscall.RawConn, iovecs *[]syscall.Iovec) (n int64, err error) {
rawReadErr := rawConn.Read(func(fd uintptr) (done bool) {
// iovecs, n and err must not be shadowed!
// NOTE: unsafe.Pointer safety rules
// https://tip.golang.org/pkg/unsafe/#Pointer
//
// (4) Conversion of a Pointer to a uintptr when calling syscall.Syscall.
// ...
// uintptr() conversions must appear within the syscall.Syscall argument list.
// (even though we are not the escape analysis Likely not )
thisReadN, _, errno := syscall.Syscall(
syscall.SYS_READV,
fd,
uintptr(unsafe.Pointer(&(*iovecs)[0])),
uintptr(len(*iovecs)),
)
if thisReadN == ^uintptr(0) {
if errno == syscall.EAGAIN {
return false
}
err = syscall.Errno(errno)
return true
}
if int(thisReadN) < 0 {
panic("unexpected return value")
}
n += int64(thisReadN) // TODO check overflow
// shift iovecs forward
for left := int(thisReadN); left > 0; {
// conversion to uint does not change value, see TestIovecLenFieldIsMachineUint, and left > 0
thisIovecConsumedCompletely := uint((*iovecs)[0].Len) <= uint(left)
if thisIovecConsumedCompletely {
// Update left, cannot go below 0 due to
// a) definition of thisIovecConsumedCompletely
// b) left > 0 due to loop invariant
// Converting .Len to int64 is thus also safe now, because it is < left < INT_MAX
left -= int((*iovecs)[0].Len)
*iovecs = (*iovecs)[1:]
} else {
// trim this iovec to remaining length
// NOTE: unsafe.Pointer safety rules
// https://tip.golang.org/pkg/unsafe/#Pointer
// (3) Conversion of a Pointer to a uintptr and back, with arithmetic.
// ...
// Note that both conversions must appear in the same expression,
// with only the intervening arithmetic between them:
(*iovecs)[0].Base = (*byte)(unsafe.Pointer(uintptr(unsafe.Pointer((*iovecs)[0].Base)) + uintptr(left)))
curVecNewLength := uint((*iovecs)[0].Len) - uint(left) // casts to uint do not change value
(*iovecs)[0].SetLen(int(curVecNewLength)) // int and uint have the same size, no change of value
break // inner
}
}
if thisReadN == 0 {
err = io.EOF
return true
}
return true
})
if rawReadErr != nil {
err = rawReadErr
}
return n, err
}
@@ -0,0 +1,149 @@
// Package grpcclientidentity makes the client identity
// provided by github.com/zrepl/zrepl/daemon/transport/serve.{AuthenticatedListener,AuthConn}
// available to gRPC service handlers.
//
// This goal is achieved through the combination of custom gRPC transport credentials and two interceptors
// (i.e. middleware).
//
// For gRPC clients, the TransportCredentials + Dialer can be used to construct a gRPC client (grpc.ClientConn)
// that uses a github.com/zrepl/zrepl/daemon/transport/connect.Connecter to connect to a server.
//
// The adaptors exposed by this package must be used together, and panic if they are not.
// See package grpchelper for a more restrictive but safe example on how the adaptors should be composed.
package grpcclientidentity
import (
"context"
"fmt"
"net"
"google.golang.org/grpc"
"google.golang.org/grpc/credentials"
"google.golang.org/grpc/peer"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/transport"
)
type Logger = logger.Logger
type GRPCDialFunction = func(context.Context, string) (net.Conn, error)
func NewDialer(logger Logger, connecter transport.Connecter) GRPCDialFunction {
return func(ctx context.Context, s string) (conn net.Conn, e error) {
nc, err := connecter.Connect(ctx)
// TODO find better place (callback from gRPC?) where to log errors
// we want the users to know, though
if err != nil {
logger.WithError(err).Error("cannot connect")
}
return nc, err
}
}
type authConnAuthType struct {
clientIdentity string
}
func (authConnAuthType) AuthType() string {
return "AuthConn"
}
type connecterAuthType struct{}
func (connecterAuthType) AuthType() string {
return "connecter"
}
type transportCredentials struct {
logger Logger
}
// Use on both sides as ServerOption or ClientOption.
func NewTransportCredentials(log Logger) credentials.TransportCredentials {
if log == nil {
log = logger.NewNullLogger()
}
return &transportCredentials{log}
}
func (c *transportCredentials) ClientHandshake(ctx context.Context, s string, rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
c.logger.WithField("url", s).WithField("connType", fmt.Sprintf("%T", rawConn)).Debug("ClientHandshake")
// do nothing, client credential is only for WithInsecure warning to go away
// the authentication is done by the connecter
return rawConn, &connecterAuthType{}, nil
}
func (c *transportCredentials) ServerHandshake(rawConn net.Conn) (net.Conn, credentials.AuthInfo, error) {
c.logger.WithField("connType", fmt.Sprintf("%T", rawConn)).Debug("ServerHandshake")
authConn, ok := rawConn.(*transport.AuthConn)
if !ok {
panic(fmt.Sprintf("NewTransportCredentials must be used with a listener that returns *transport.AuthConn, got %T", rawConn))
}
return rawConn, &authConnAuthType{authConn.ClientIdentity()}, nil
}
func (*transportCredentials) Info() credentials.ProtocolInfo {
return credentials.ProtocolInfo{} // TODO
}
func (t *transportCredentials) Clone() credentials.TransportCredentials {
var x = *t
return &x
}
func (*transportCredentials) OverrideServerName(string) error {
panic("not implemented")
}
type ContextInterceptorData interface {
FullMethod() string
ClientIdentity() string
}
type contextInterceptorData struct {
fullMethod string
clientIdentity string
}
func (d contextInterceptorData) FullMethod() string { return d.fullMethod }
func (d contextInterceptorData) ClientIdentity() string { return d.clientIdentity }
type Interceptor = func(ctx context.Context, data ContextInterceptorData, handler func(ctx context.Context))
func NewInterceptors(logger Logger, clientIdentityKey interface{}, interceptor Interceptor) (unary grpc.UnaryServerInterceptor, stream grpc.StreamServerInterceptor) {
unary = func(ctx context.Context, req interface{}, info *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (interface{}, error) {
logger.WithField("fullMethod", info.FullMethod).Debug("request")
p, ok := peer.FromContext(ctx)
if !ok {
panic("peer.FromContext expected to return a peer in grpc.UnaryServerInterceptor")
}
peerAddr := ""
if p.Addr != nil { // https://github.com/zrepl/zrepl/issues/598
peerAddr = p.Addr.String()
}
logger.WithField("peer_addr", peerAddr).Debug("peer addr")
a, ok := p.AuthInfo.(*authConnAuthType)
if !ok {
panic(fmt.Sprintf("NewInterceptors must be used in combination with grpc.NewTransportCredentials, but got auth type %T", p.AuthInfo))
}
logger.WithField("peer_client_identity", a.clientIdentity).Debug("peer client identity")
ctx = context.WithValue(ctx, clientIdentityKey, a.clientIdentity)
data := contextInterceptorData{
fullMethod: info.FullMethod,
clientIdentity: a.clientIdentity,
}
var (
resp interface{}
err error
)
interceptor(ctx, data, func(ctx context.Context) {
resp, err = handler(ctx, req) // no-shadow
})
return resp, err
}
stream = func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error {
panic("unimplemented")
}
return
}
@@ -0,0 +1,16 @@
syntax = "proto3";
option go_package = ".;pdu";
package pdu;
service Greeter {
rpc Greet(GreetRequest) returns (GreetResponse) {}
}
message GreetRequest {
string name = 1;
}
message GreetResponse {
string msg = 1;
}
@@ -0,0 +1,114 @@
// This package demonstrates how the grpcclientidentity package can be used
// to set up a gRPC greeter service.
package main
import (
"context"
"flag"
"fmt"
"os"
"time"
"github.com/zrepl/zrepl/internal/config"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/rpc/grpcclientidentity/example/pdu"
"github.com/zrepl/zrepl/internal/rpc/grpcclientidentity/grpchelper"
"github.com/zrepl/zrepl/internal/transport/tcp"
)
var args struct {
mode string
}
var log = logger.NewStderrDebugLogger()
func main() {
flag.StringVar(&args.mode, "mode", "", "client|server")
flag.Parse()
switch args.mode {
case "client":
client()
case "server":
server()
default:
log.Printf("unknown mode %q")
os.Exit(1)
}
}
func onErr(err error, format string, args ...interface{}) {
log.WithError(err).Error(fmt.Sprintf("%s: %s", fmt.Sprintf(format, args...), err))
os.Exit(1)
}
func client() {
cn, err := tcp.TCPConnecterFromConfig(&config.TCPConnect{
ConnectCommon: config.ConnectCommon{
Type: "tcp",
},
Address: "127.0.0.1:8080",
DialTimeout: 10 * time.Second,
})
if err != nil {
onErr(err, "build connecter error")
}
clientConn := grpchelper.ClientConn(cn, log)
defer clientConn.Close()
// normal usage from here on
client := pdu.NewGreeterClient(clientConn)
resp, err := client.Greet(context.Background(), &pdu.GreetRequest{Name: "somethingimadeup"})
if err != nil {
onErr(err, "RPC error")
}
fmt.Printf("got response:\n\t%s\n", resp.GetMsg())
}
const clientIdentityKey = "clientIdentity"
func server() {
authListenerFactory, err := tcp.TCPListenerFactoryFromConfig(nil, &config.TCPServe{
ServeCommon: config.ServeCommon{
Type: "tcp",
},
Listen: "127.0.0.1:8080",
Clients: map[string]string{
"127.0.0.1": "localclient",
"::1": "localclient",
},
})
if err != nil {
onErr(err, "cannot build listener factory")
}
log := logger.NewStderrDebugLogger()
authListener, err := authListenerFactory()
if err != nil {
onErr(err, "cannot listen")
}
srv, serve := grpchelper.NewServer(authListener, clientIdentityKey, log, nil)
svc := &greeter{prepend: "hello "}
pdu.RegisterGreeterServer(srv, svc)
if err := serve(); err != nil {
onErr(err, "error serving")
}
}
type greeter struct {
pdu.UnsafeGreeterServer
prepend string
}
func (g *greeter) Greet(ctx context.Context, r *pdu.GreetRequest) (*pdu.GreetResponse, error) {
ci, _ := ctx.Value(clientIdentityKey).(string)
log.WithField("clientIdentity", ci).Info("Greet() request") // show that we got the client identity
return &pdu.GreetResponse{Msg: fmt.Sprintf("%s%s (clientIdentity=%q)", g.prepend, r.GetName(), ci)}, nil
}
@@ -0,0 +1,209 @@
// Code generated by protoc-gen-go. DO NOT EDIT.
// versions:
// protoc-gen-go v1.34.2
// protoc v5.28.0
// source: grpcauth.proto
package pdu
import (
protoreflect "google.golang.org/protobuf/reflect/protoreflect"
protoimpl "google.golang.org/protobuf/runtime/protoimpl"
reflect "reflect"
sync "sync"
)
const (
// Verify that this generated code is sufficiently up-to-date.
_ = protoimpl.EnforceVersion(20 - protoimpl.MinVersion)
// Verify that runtime/protoimpl is sufficiently up-to-date.
_ = protoimpl.EnforceVersion(protoimpl.MaxVersion - 20)
)
type GreetRequest struct {
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Name string `protobuf:"bytes,1,opt,name=name,proto3" json:"name,omitempty"`
}
func (x *GreetRequest) Reset() {
*x = GreetRequest{}
if protoimpl.UnsafeEnabled {
mi := &file_grpcauth_proto_msgTypes[0]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *GreetRequest) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*GreetRequest) ProtoMessage() {}
func (x *GreetRequest) ProtoReflect() protoreflect.Message {
mi := &file_grpcauth_proto_msgTypes[0]
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use GreetRequest.ProtoReflect.Descriptor instead.
func (*GreetRequest) Descriptor() ([]byte, []int) {
return file_grpcauth_proto_rawDescGZIP(), []int{0}
}
func (x *GreetRequest) GetName() string {
if x != nil {
return x.Name
}
return ""
}
type GreetResponse struct {
state protoimpl.MessageState
sizeCache protoimpl.SizeCache
unknownFields protoimpl.UnknownFields
Msg string `protobuf:"bytes,1,opt,name=msg,proto3" json:"msg,omitempty"`
}
func (x *GreetResponse) Reset() {
*x = GreetResponse{}
if protoimpl.UnsafeEnabled {
mi := &file_grpcauth_proto_msgTypes[1]
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
ms.StoreMessageInfo(mi)
}
}
func (x *GreetResponse) String() string {
return protoimpl.X.MessageStringOf(x)
}
func (*GreetResponse) ProtoMessage() {}
func (x *GreetResponse) ProtoReflect() protoreflect.Message {
mi := &file_grpcauth_proto_msgTypes[1]
if protoimpl.UnsafeEnabled && x != nil {
ms := protoimpl.X.MessageStateOf(protoimpl.Pointer(x))
if ms.LoadMessageInfo() == nil {
ms.StoreMessageInfo(mi)
}
return ms
}
return mi.MessageOf(x)
}
// Deprecated: Use GreetResponse.ProtoReflect.Descriptor instead.
func (*GreetResponse) Descriptor() ([]byte, []int) {
return file_grpcauth_proto_rawDescGZIP(), []int{1}
}
func (x *GreetResponse) GetMsg() string {
if x != nil {
return x.Msg
}
return ""
}
var File_grpcauth_proto protoreflect.FileDescriptor
var file_grpcauth_proto_rawDesc = []byte{
0x0a, 0x0e, 0x67, 0x72, 0x70, 0x63, 0x61, 0x75, 0x74, 0x68, 0x2e, 0x70, 0x72, 0x6f, 0x74, 0x6f,
0x12, 0x03, 0x70, 0x64, 0x75, 0x22, 0x22, 0x0a, 0x0c, 0x47, 0x72, 0x65, 0x65, 0x74, 0x52, 0x65,
0x71, 0x75, 0x65, 0x73, 0x74, 0x12, 0x12, 0x0a, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x18, 0x01, 0x20,
0x01, 0x28, 0x09, 0x52, 0x04, 0x6e, 0x61, 0x6d, 0x65, 0x22, 0x21, 0x0a, 0x0d, 0x47, 0x72, 0x65,
0x65, 0x74, 0x52, 0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x12, 0x10, 0x0a, 0x03, 0x6d, 0x73,
0x67, 0x18, 0x01, 0x20, 0x01, 0x28, 0x09, 0x52, 0x03, 0x6d, 0x73, 0x67, 0x32, 0x3b, 0x0a, 0x07,
0x47, 0x72, 0x65, 0x65, 0x74, 0x65, 0x72, 0x12, 0x30, 0x0a, 0x05, 0x47, 0x72, 0x65, 0x65, 0x74,
0x12, 0x11, 0x2e, 0x70, 0x64, 0x75, 0x2e, 0x47, 0x72, 0x65, 0x65, 0x74, 0x52, 0x65, 0x71, 0x75,
0x65, 0x73, 0x74, 0x1a, 0x12, 0x2e, 0x70, 0x64, 0x75, 0x2e, 0x47, 0x72, 0x65, 0x65, 0x74, 0x52,
0x65, 0x73, 0x70, 0x6f, 0x6e, 0x73, 0x65, 0x22, 0x00, 0x42, 0x07, 0x5a, 0x05, 0x2e, 0x3b, 0x70,
0x64, 0x75, 0x62, 0x06, 0x70, 0x72, 0x6f, 0x74, 0x6f, 0x33,
}
var (
file_grpcauth_proto_rawDescOnce sync.Once
file_grpcauth_proto_rawDescData = file_grpcauth_proto_rawDesc
)
func file_grpcauth_proto_rawDescGZIP() []byte {
file_grpcauth_proto_rawDescOnce.Do(func() {
file_grpcauth_proto_rawDescData = protoimpl.X.CompressGZIP(file_grpcauth_proto_rawDescData)
})
return file_grpcauth_proto_rawDescData
}
var file_grpcauth_proto_msgTypes = make([]protoimpl.MessageInfo, 2)
var file_grpcauth_proto_goTypes = []any{
(*GreetRequest)(nil), // 0: pdu.GreetRequest
(*GreetResponse)(nil), // 1: pdu.GreetResponse
}
var file_grpcauth_proto_depIdxs = []int32{
0, // 0: pdu.Greeter.Greet:input_type -> pdu.GreetRequest
1, // 1: pdu.Greeter.Greet:output_type -> pdu.GreetResponse
1, // [1:2] is the sub-list for method output_type
0, // [0:1] is the sub-list for method input_type
0, // [0:0] is the sub-list for extension type_name
0, // [0:0] is the sub-list for extension extendee
0, // [0:0] is the sub-list for field type_name
}
func init() { file_grpcauth_proto_init() }
func file_grpcauth_proto_init() {
if File_grpcauth_proto != nil {
return
}
if !protoimpl.UnsafeEnabled {
file_grpcauth_proto_msgTypes[0].Exporter = func(v any, i int) any {
switch v := v.(*GreetRequest); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
file_grpcauth_proto_msgTypes[1].Exporter = func(v any, i int) any {
switch v := v.(*GreetResponse); i {
case 0:
return &v.state
case 1:
return &v.sizeCache
case 2:
return &v.unknownFields
default:
return nil
}
}
}
type x struct{}
out := protoimpl.TypeBuilder{
File: protoimpl.DescBuilder{
GoPackagePath: reflect.TypeOf(x{}).PkgPath(),
RawDescriptor: file_grpcauth_proto_rawDesc,
NumEnums: 0,
NumMessages: 2,
NumExtensions: 0,
NumServices: 1,
},
GoTypes: file_grpcauth_proto_goTypes,
DependencyIndexes: file_grpcauth_proto_depIdxs,
MessageInfos: file_grpcauth_proto_msgTypes,
}.Build()
File_grpcauth_proto = out.File
file_grpcauth_proto_rawDesc = nil
file_grpcauth_proto_goTypes = nil
file_grpcauth_proto_depIdxs = nil
}
@@ -0,0 +1,121 @@
// Code generated by protoc-gen-go-grpc. DO NOT EDIT.
// versions:
// - protoc-gen-go-grpc v1.5.1
// - protoc v5.28.0
// source: grpcauth.proto
package pdu
import (
context "context"
grpc "google.golang.org/grpc"
codes "google.golang.org/grpc/codes"
status "google.golang.org/grpc/status"
)
// This is a compile-time assertion to ensure that this generated file
// is compatible with the grpc package it is being compiled against.
// Requires gRPC-Go v1.64.0 or later.
const _ = grpc.SupportPackageIsVersion9
const (
Greeter_Greet_FullMethodName = "/pdu.Greeter/Greet"
)
// GreeterClient is the client API for Greeter service.
//
// For semantics around ctx use and closing/ending streaming RPCs, please refer to https://pkg.go.dev/google.golang.org/grpc/?tab=doc#ClientConn.NewStream.
type GreeterClient interface {
Greet(ctx context.Context, in *GreetRequest, opts ...grpc.CallOption) (*GreetResponse, error)
}
type greeterClient struct {
cc grpc.ClientConnInterface
}
func NewGreeterClient(cc grpc.ClientConnInterface) GreeterClient {
return &greeterClient{cc}
}
func (c *greeterClient) Greet(ctx context.Context, in *GreetRequest, opts ...grpc.CallOption) (*GreetResponse, error) {
cOpts := append([]grpc.CallOption{grpc.StaticMethod()}, opts...)
out := new(GreetResponse)
err := c.cc.Invoke(ctx, Greeter_Greet_FullMethodName, in, out, cOpts...)
if err != nil {
return nil, err
}
return out, nil
}
// GreeterServer is the server API for Greeter service.
// All implementations must embed UnimplementedGreeterServer
// for forward compatibility.
type GreeterServer interface {
Greet(context.Context, *GreetRequest) (*GreetResponse, error)
mustEmbedUnimplementedGreeterServer()
}
// UnimplementedGreeterServer must be embedded to have
// forward compatible implementations.
//
// NOTE: this should be embedded by value instead of pointer to avoid a nil
// pointer dereference when methods are called.
type UnimplementedGreeterServer struct{}
func (UnimplementedGreeterServer) Greet(context.Context, *GreetRequest) (*GreetResponse, error) {
return nil, status.Errorf(codes.Unimplemented, "method Greet not implemented")
}
func (UnimplementedGreeterServer) mustEmbedUnimplementedGreeterServer() {}
func (UnimplementedGreeterServer) testEmbeddedByValue() {}
// UnsafeGreeterServer may be embedded to opt out of forward compatibility for this service.
// Use of this interface is not recommended, as added methods to GreeterServer will
// result in compilation errors.
type UnsafeGreeterServer interface {
mustEmbedUnimplementedGreeterServer()
}
func RegisterGreeterServer(s grpc.ServiceRegistrar, srv GreeterServer) {
// If the following call pancis, it indicates UnimplementedGreeterServer was
// embedded by pointer and is nil. This will cause panics if an
// unimplemented method is ever invoked, so we test this at initialization
// time to prevent it from happening at runtime later due to I/O.
if t, ok := srv.(interface{ testEmbeddedByValue() }); ok {
t.testEmbeddedByValue()
}
s.RegisterService(&Greeter_ServiceDesc, srv)
}
func _Greeter_Greet_Handler(srv interface{}, ctx context.Context, dec func(interface{}) error, interceptor grpc.UnaryServerInterceptor) (interface{}, error) {
in := new(GreetRequest)
if err := dec(in); err != nil {
return nil, err
}
if interceptor == nil {
return srv.(GreeterServer).Greet(ctx, in)
}
info := &grpc.UnaryServerInfo{
Server: srv,
FullMethod: Greeter_Greet_FullMethodName,
}
handler := func(ctx context.Context, req interface{}) (interface{}, error) {
return srv.(GreeterServer).Greet(ctx, req.(*GreetRequest))
}
return interceptor(ctx, in, info, handler)
}
// Greeter_ServiceDesc is the grpc.ServiceDesc for Greeter service.
// It's only intended for direct use with grpc.RegisterService,
// and not to be introspected or modified (even as a copy)
var Greeter_ServiceDesc = grpc.ServiceDesc{
ServiceName: "pdu.Greeter",
HandlerType: (*GreeterServer)(nil),
Methods: []grpc.MethodDesc{
{
MethodName: "Greet",
Handler: _Greeter_Greet_Handler,
},
},
Streams: []grpc.StreamDesc{},
Metadata: "grpcauth.proto",
}
@@ -0,0 +1,75 @@
// Package grpchelper wraps the adaptors implemented by package grpcclientidentity into a less flexible API
// which, however, ensures that the individual adaptor primitive's expectations are met and hence do not panic.
package grpchelper
import (
"time"
"google.golang.org/grpc"
"google.golang.org/grpc/keepalive"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/rpc/grpcclientidentity"
"github.com/zrepl/zrepl/internal/rpc/netadaptor"
"github.com/zrepl/zrepl/internal/transport"
)
// The following constants are relevant for interoperability.
// We use the same values for client & server, because zrepl is more
// symmetrical ("one source, one sink") instead of the typical
// gRPC scenario ("many clients, single server")
const (
StartKeepalivesAfterInactivityDuration = 5 * time.Second
KeepalivePeerTimeout = 10 * time.Second
)
type Logger = logger.Logger
// ClientConn is an easy-to-use wrapper around the Dialer and TransportCredentials interface
// to produce a grpc.ClientConn
func ClientConn(cn transport.Connecter, log Logger) *grpc.ClientConn {
ka := grpc.WithKeepaliveParams(keepalive.ClientParameters{
Time: StartKeepalivesAfterInactivityDuration,
Timeout: KeepalivePeerTimeout,
PermitWithoutStream: true,
})
dialerOption := grpc.WithContextDialer(grpcclientidentity.NewDialer(log, cn))
cred := grpc.WithTransportCredentials(grpcclientidentity.NewTransportCredentials(log))
// we use context.Background without a timeout here because we don't set grpc.WithBlock
// => docs: "In the non-blocking case, the ctx does not act against the connection. It only controls the setup steps."
cc, err := grpc.NewClient("passthrough://doesntmatterdonebydialer", dialerOption, cred, ka)
if err != nil {
log.WithError(err).Error("cannot create gRPC client conn (non-blocking)")
// It's ok to panic here: the we call grpc.DialContext without the
// (grpc.WithBlock) dial option, and at the time of writing, the grpc
// docs state that no connection attempt is made in that case.
// Hence, any error that occurs is due to DialOptions or similar,
// and thus indicative of an implementation error.
panic(err)
}
return cc
}
// NewServer is a convenience interface around the TransportCredentials and Interceptors interface.
func NewServer(authListener transport.AuthenticatedListener, clientIdentityKey interface{}, logger grpcclientidentity.Logger, ctxInterceptor grpcclientidentity.Interceptor) (srv *grpc.Server, serve func() error) {
ka := grpc.KeepaliveParams(keepalive.ServerParameters{
Time: StartKeepalivesAfterInactivityDuration,
Timeout: KeepalivePeerTimeout,
})
ep := grpc.KeepaliveEnforcementPolicy(keepalive.EnforcementPolicy{
MinTime: StartKeepalivesAfterInactivityDuration / 2, // avoid skew
PermitWithoutStream: true,
})
tcs := grpcclientidentity.NewTransportCredentials(logger)
unary, stream := grpcclientidentity.NewInterceptors(logger, clientIdentityKey, ctxInterceptor)
srv = grpc.NewServer(grpc.Creds(tcs), grpc.UnaryInterceptor(unary), grpc.StreamInterceptor(stream), ka, ep)
serve = func() error {
if err := srv.Serve(netadaptor.New(authListener, logger)); err != nil {
return err
}
return nil
}
return srv, serve
}
@@ -0,0 +1,87 @@
// Package netadaptor implements an adaptor from
// transport.AuthenticatedListener to net.Listener.
package netadaptor
import (
"context"
"fmt"
"net"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/transport"
)
type Logger = logger.Logger
type acceptRes struct {
conn *transport.AuthConn
err error
}
type acceptReq struct {
callback chan acceptRes
}
type Listener struct {
al transport.AuthenticatedListener
logger Logger
accepts chan acceptReq
stop chan struct{}
}
// Consume the given authListener and wrap it as a *Listener, which implements net.Listener.
// The returned net.Listener must be closed.
// The wrapped authListener is closed when the returned net.Listener is closed.
func New(authListener transport.AuthenticatedListener, l Logger) *Listener {
if l == nil {
l = logger.NewNullLogger()
}
a := &Listener{
al: authListener,
logger: l,
accepts: make(chan acceptReq),
stop: make(chan struct{}),
}
go a.handleAccept()
return a
}
// The returned net.Conn is guaranteed to be *transport.AuthConn, i.e., the type of connection
// returned by the wrapped transport.AuthenticatedListener.
func (a Listener) Accept() (net.Conn, error) {
req := acceptReq{make(chan acceptRes, 1)}
select {
case a.accepts <- req:
case <-a.stop:
return nil, fmt.Errorf("already closed") // TODO net.Error
}
res, ok := <-req.callback
if !ok {
return nil, fmt.Errorf("already closed") // TODO net.Error
}
return res.conn, res.err
}
func (a Listener) handleAccept() {
for {
select {
case <-a.stop:
a.logger.Debug("handleAccept stop accepting")
return
case req := <-a.accepts:
a.logger.Debug("accept authListener")
authConn, err := a.al.Accept(context.Background())
req.callback <- acceptRes{authConn, err}
}
}
}
func (a Listener) Addr() net.Addr {
return a.al.Addr()
}
func (a Listener) Close() error {
close(a.stop)
return a.al.Close()
}
+224
View File
@@ -0,0 +1,224 @@
package rpc
import (
"context"
"errors"
"fmt"
"io"
"net"
"sync"
"sync/atomic"
"time"
"google.golang.org/grpc"
"github.com/zrepl/zrepl/internal/daemon/logging/trace"
"github.com/google/uuid"
"github.com/zrepl/zrepl/internal/replication/logic"
"github.com/zrepl/zrepl/internal/replication/logic/pdu"
"github.com/zrepl/zrepl/internal/rpc/dataconn"
"github.com/zrepl/zrepl/internal/rpc/grpcclientidentity/grpchelper"
"github.com/zrepl/zrepl/internal/rpc/versionhandshake"
"github.com/zrepl/zrepl/internal/transport"
"github.com/zrepl/zrepl/internal/util/envconst"
)
// Client implements the active side of a replication setup.
// It satisfies the Endpoint, Sender and Receiver interface defined by package replication.
type Client struct {
dataClient *dataconn.Client
controlClient pdu.ReplicationClient // this the grpc client instance, see constructor
controlConn *grpc.ClientConn
loggers Loggers
closed chan struct{}
}
var _ logic.Endpoint = &Client{}
var _ logic.Sender = &Client{}
var _ logic.Receiver = &Client{}
type DialContextFunc = func(ctx context.Context, network string, addr string) (net.Conn, error)
// config must be validated, NewClient will panic if it is not valid
func NewClient(cn transport.Connecter, loggers Loggers) *Client {
cn = versionhandshake.Connecter(cn, envconst.Duration("ZREPL_RPC_CLIENT_VERSIONHANDSHAKE_TIMEOUT", 10*time.Second))
muxedConnecter := mux(cn)
c := &Client{
loggers: loggers,
closed: make(chan struct{}),
}
grpcConn := grpchelper.ClientConn(muxedConnecter.control, loggers.Control)
go func() {
ctx, cancel := context.WithCancel(context.Background())
go func() {
<-c.closed
cancel()
}()
defer cancel()
for ctx.Err() == nil {
state := grpcConn.GetState()
loggers.General.WithField("grpc_state", state.String()).Debug("grpc state change")
grpcConn.WaitForStateChange(ctx, state)
}
}()
c.controlClient = pdu.NewReplicationClient(grpcConn)
c.controlConn = grpcConn
c.dataClient = dataconn.NewClient(muxedConnecter.data, loggers.Data)
return c
}
func (c *Client) Close() {
close(c.closed)
if err := c.controlConn.Close(); err != nil {
c.loggers.General.WithError(err).Error("cannot close control connection")
}
// TODO c.dataClient should have Close()
}
// callers must ensure that the returned io.ReadCloser is closed
// TODO expose dataClient interface to the outside world
func (c *Client) Send(ctx context.Context, r *pdu.SendReq) (*pdu.SendRes, io.ReadCloser, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.Send")
defer endSpan()
// TODO the returned sendStream may return a read error created by the remote side
res, stream, err := c.dataClient.ReqSend(ctx, r)
if err != nil {
return nil, nil, err
}
if stream == nil {
return res, nil, nil
}
return res, stream, nil
}
func (c *Client) Receive(ctx context.Context, req *pdu.ReceiveReq, stream io.ReadCloser) (*pdu.ReceiveRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.Receive")
defer endSpan()
return c.dataClient.ReqRecv(ctx, req, stream)
}
func (c *Client) SendDry(ctx context.Context, in *pdu.SendReq) (*pdu.SendRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.SendDry")
defer endSpan()
return c.controlClient.SendDry(ctx, in)
}
func (c *Client) ListFilesystems(ctx context.Context, in *pdu.ListFilesystemReq) (*pdu.ListFilesystemRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.ListFilesystems")
defer endSpan()
return c.controlClient.ListFilesystems(ctx, in)
}
func (c *Client) ListFilesystemVersions(ctx context.Context, in *pdu.ListFilesystemVersionsReq) (*pdu.ListFilesystemVersionsRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.ListFilesystemVersions")
defer endSpan()
return c.controlClient.ListFilesystemVersions(ctx, in)
}
func (c *Client) DestroySnapshots(ctx context.Context, in *pdu.DestroySnapshotsReq) (*pdu.DestroySnapshotsRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.DestroySnapshots")
defer endSpan()
return c.controlClient.DestroySnapshots(ctx, in)
}
func (c *Client) ReplicationCursor(ctx context.Context, in *pdu.ReplicationCursorReq) (*pdu.ReplicationCursorRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.ReplicationCursor")
defer endSpan()
return c.controlClient.ReplicationCursor(ctx, in)
}
func (c *Client) SendCompleted(ctx context.Context, in *pdu.SendCompletedReq) (*pdu.SendCompletedRes, error) {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.SendCompleted")
defer endSpan()
return c.controlClient.SendCompleted(ctx, in)
}
func (c *Client) WaitForConnectivity(ctx context.Context) error {
ctx, endSpan := trace.WithSpan(ctx, "rpc.client.WaitForConnectivity")
defer endSpan()
ctx, cancel := context.WithCancel(ctx)
defer cancel()
msg := uuid.New().String()
req := pdu.PingReq{Message: msg}
var ctrlOk, dataOk int32
loggers := GetLoggersOrPanic(ctx)
var wg sync.WaitGroup
wg.Add(2)
checkRes := func(res *pdu.PingRes, err error, logger Logger, okVar *int32) {
if err == nil && res.GetEcho() != req.GetMessage() {
err = errors.New("pilot message not echoed correctly")
}
if err == context.Canceled {
err = nil
}
if err != nil {
logger.WithError(err).Error("ping failed")
atomic.StoreInt32(okVar, 0)
cancel()
} else {
atomic.StoreInt32(okVar, 1)
}
}
go func() {
defer wg.Done()
ctrl, ctrlErr := c.controlClient.Ping(ctx, &req, grpc.WaitForReady(true))
checkRes(ctrl, ctrlErr, loggers.Control, &ctrlOk)
}()
go func() {
defer wg.Done()
for ctx.Err() == nil {
data, dataErr := c.dataClient.ReqPing(ctx, &req)
// dataClient uses transport.Connecter, which doesn't expose WaitForReady(true)
// => we need to mask dial timeouts
if err, ok := dataErr.(interface{ Timeout() bool }); ok && err.Timeout() {
// Rate-limit pings here in case Timeout() == true is a mis-classification
// or returns immediately (this is a tight loop in that case)
// TODO keep this in lockstep with controlClient
// => don't use FailFast for control, but check that both control and data worked
time.Sleep(envconst.Duration("ZREPL_RPC_DATACONN_PING_SLEEP", 1*time.Second))
continue
}
// it's not a dial timeout,
checkRes(data, dataErr, loggers.Data, &dataOk)
return
}
}()
wg.Wait()
var what string
if ctrlOk == 1 && dataOk == 1 {
return nil
}
if ctrlOk == 0 {
what += "control"
}
if dataOk == 0 {
if len(what) > 0 {
what += " and data"
} else {
what += "data"
}
}
return fmt.Errorf("%s rpc failed to respond to ping rpcs", what)
}
func (c *Client) ResetConnectBackoff() {
c.controlConn.ResetConnectBackoff()
}
+22
View File
@@ -0,0 +1,22 @@
package rpc
import (
"fmt"
"os"
)
//nolint:unused
var debugEnabled bool = false
func init() {
if os.Getenv("ZREPL_RPC_DEBUG") != "" {
debugEnabled = true
}
}
//nolint:deadcode,unused
func debug(format string, args ...interface{}) {
if debugEnabled {
fmt.Fprintf(os.Stderr, "rpc: %s\n", fmt.Sprintf(format, args...))
}
}
+115
View File
@@ -0,0 +1,115 @@
// Package rpc implements zrepl daemon-to-daemon RPC protocol
// on top of a transport provided by package transport.
// The zrepl documentation refers to the client as the
// `active side` and to the server as the `passive side`.
//
// # Design Considerations
//
// zrepl has non-standard requirements to remote procedure calls (RPC):
// whereas the coordination of replication (the planning phase) mostly
// consists of regular unary calls, the actual data transfer requires
// a high-throughput, low-overhead solution.
//
// Specifically, the requirements for data transfer is to perform
// a copy of an io.Reader over the wire, such that an io.EOF of the original
// reader corresponds to an io.EOF on the receiving side.
// If any other error occurs while reading from the original io.Reader
// on the sender, the receiver should receive the contents of that error
// in some form (usually as a trailer message)
// A common implementation technique for above data transfer is chunking,
// for example in HTTP:
// https://tools.ietf.org/html/rfc2616#section-3.6.1
//
// With regards to regular RPCs, gRPC is a popular implementation based
// on protocol buffers and thus code generation.
// gRPC also supports bi-directional streaming RPC, and it is possible
// to implement chunked transfer through the use of streams.
//
// For zrepl however, both HTTP and manually implemented chunked transfer
// using gRPC were found to have significant CPU overhead at transfer
// speeds to be expected even with hobbyist users.
//
// However, it is nice to piggyback on the code generation provided
// by protobuf / gRPC, in particular since the majority of call types
// are regular unary RPCs for which the higher overhead of gRPC is acceptable.
//
// Hence, this package attempts to combine the best of both worlds:
//
// # GRPC for Coordination and Dataconn for Bulk Data Transfer
//
// This package's Client uses its transport.Connecter to maintain
// separate control and data connections to the Server.
// The control connection is used by an instance of pdu.ReplicationClient
// whose 'regular' unary RPC calls are re-exported.
// The data connection is used by an instance of dataconn.Client and
// is used for bulk data transfer, namely `Send` and `Receive`.
//
// The following ASCII diagram gives an overview of how the individual
// building blocks are glued together:
//
// +------------+
// | rpc.Client |
// +------------+
// | |
// +--------+ +------------+
// | |
// +---------v-----------+ +--------v------+
// |pdu.ReplicationClient| |dataconn.Client|
// +---------------------+ +--------v------+
// | label: label: |
// | zrepl_control zrepl_data |
// +--------+ +------------+
// | |
// +--v---------v---+
// | transportmux |
// +-------+--------+
// | uses
// +-------v--------+
// |versionhandshake|
// +-------+--------+
// | uses
// +------v------+
// | transport |
// +------+------+
// |
// NETWORK
// |
// +------+------+
// | transport |
// +------^------+
// | uses
// +-------+--------+
// |versionhandshake|
// +-------^--------+
// | uses
// +-------+--------+
// | transportmux |
// +--^--------^----+
// | |
// +--------+ --------------+ ---
// | | |
// | label: label: | |
// | zrepl_control zrepl_data | |
// +-----+----+ +-----------+---+ |
// |netadaptor| |dataconn.Server| | rpc.Server
// | + | +------+--------+ |
// |grpcclient| | |
// |identity | | |
// +-----+----+ | |
// | | |
// +---------v-----------+ | |
// |pdu.ReplicationServer| | |
// +---------+-----------+ | |
// | | ---
// +----------+ +------------+
// | |
// +---v--v-----+
// | Handler |
// +------------+
// (usually endpoint.{Sender,Receiver})
package rpc
// edit trick for the ASCII art above:
// - remove the comments //
// - vim: set virtualedit+=all
// - vim: set ft=text
+25
View File
@@ -0,0 +1,25 @@
package rpc
import (
"context"
"github.com/zrepl/zrepl/internal/daemon/logging"
"github.com/zrepl/zrepl/internal/logger"
)
type Logger = logger.Logger
// All fields must be non-nil
type Loggers struct {
General Logger
Control Logger
Data Logger
}
func GetLoggersOrPanic(ctx context.Context) Loggers {
return Loggers{
General: logging.GetLogger(ctx, logging.SubsysRPC),
Control: logging.GetLogger(ctx, logging.SubsysRPCControl),
Data: logging.GetLogger(ctx, logging.SubsysRPCData),
}
}
+57
View File
@@ -0,0 +1,57 @@
package rpc
import (
"context"
"time"
"github.com/zrepl/zrepl/internal/rpc/transportmux"
"github.com/zrepl/zrepl/internal/transport"
"github.com/zrepl/zrepl/internal/util/envconst"
)
type demuxedListener struct {
control, data transport.AuthenticatedListener
}
const (
transportmuxLabelControl string = "zrepl_control"
transportmuxLabelData string = "zrepl_data"
)
func demux(serveCtx context.Context, listener transport.AuthenticatedListener) demuxedListener {
listeners, err := transportmux.Demux(
serveCtx, listener,
[]string{transportmuxLabelControl, transportmuxLabelData},
envconst.Duration("ZREPL_TRANSPORT_DEMUX_TIMEOUT", 10*time.Second),
)
if err != nil {
// transportmux API guarantees that the returned error can only be due
// to invalid API usage (i.e. labels too long)
panic(err)
}
return demuxedListener{
control: listeners[transportmuxLabelControl],
data: listeners[transportmuxLabelData],
}
}
type muxedConnecter struct {
control, data transport.Connecter
}
func mux(rawConnecter transport.Connecter) muxedConnecter {
muxedConnecters, err := transportmux.MuxConnecter(
rawConnecter,
[]string{transportmuxLabelControl, transportmuxLabelData},
envconst.Duration("ZREPL_TRANSPORT_MUX_TIMEOUT", 10*time.Second),
)
if err != nil {
// transportmux API guarantees that the returned error can only be due
// to invalid API usage (i.e. labels too long)
panic(err)
}
return muxedConnecter{
control: muxedConnecters[transportmuxLabelControl],
data: muxedConnecters[transportmuxLabelData],
}
}
+132
View File
@@ -0,0 +1,132 @@
package rpc
import (
"context"
"time"
"github.com/zrepl/zrepl/internal/endpoint"
"github.com/zrepl/zrepl/internal/replication/logic/pdu"
"github.com/zrepl/zrepl/internal/rpc/dataconn"
"github.com/zrepl/zrepl/internal/rpc/grpcclientidentity"
"github.com/zrepl/zrepl/internal/rpc/grpcclientidentity/grpchelper"
"github.com/zrepl/zrepl/internal/rpc/versionhandshake"
"github.com/zrepl/zrepl/internal/transport"
"github.com/zrepl/zrepl/internal/util/envconst"
)
type Handler interface {
pdu.ReplicationServer
dataconn.Handler
}
type serveFunc func(ctx context.Context, demuxedListener transport.AuthenticatedListener, errOut chan<- error)
// Server abstracts the accept and request routing infrastructure for the
// passive side of a replication setup.
type Server struct {
logger Logger
handler Handler
controlServerServe serveFunc
dataServer *dataconn.Server
dataServerServe serveFunc
}
type HandlerContextInterceptorData interface {
FullMethod() string
ClientIdentity() string
}
type interceptorData struct {
prefixMethod string
wrapped HandlerContextInterceptorData
}
func (d interceptorData) ClientIdentity() string { return d.wrapped.ClientIdentity() }
func (d interceptorData) FullMethod() string { return d.prefixMethod + d.wrapped.FullMethod() }
type HandlerContextInterceptor func(ctx context.Context, data HandlerContextInterceptorData, handler func(ctx context.Context))
// config must be valid (use its Validate function).
func NewServer(handler Handler, loggers Loggers, ctxInterceptor HandlerContextInterceptor) *Server {
// setup control server
controlServerServe := func(ctx context.Context, controlListener transport.AuthenticatedListener, errOut chan<- error) {
var controlCtxInterceptor grpcclientidentity.Interceptor = func(ctx context.Context, data grpcclientidentity.ContextInterceptorData, handler func(ctx context.Context)) {
ctxInterceptor(ctx, interceptorData{"control://", data}, handler)
}
controlServer, serve := grpchelper.NewServer(controlListener, endpoint.ClientIdentityKey, loggers.Control, controlCtxInterceptor)
pdu.RegisterReplicationServer(controlServer, handler)
// give time for graceful stop until deadline expires, then hard stop
go func() {
<-ctx.Done()
if dl, ok := ctx.Deadline(); ok {
go time.AfterFunc(dl.Sub(dl), controlServer.Stop)
}
loggers.Control.Debug("gracefully shutting down control server")
controlServer.GracefulStop()
loggers.Control.Debug("gracdeful shut down of control server complete")
}()
errOut <- serve()
}
// setup data server
dataServerClientIdentitySetter := func(ctx context.Context, wire *transport.AuthConn) (context.Context, *transport.AuthConn) {
ci := wire.ClientIdentity()
ctx = context.WithValue(ctx, endpoint.ClientIdentityKey, ci)
return ctx, wire
}
var dataCtxInterceptor dataconn.ContextInterceptor = func(ctx context.Context, data dataconn.ContextInterceptorData, handler func(ctx context.Context)) {
ctxInterceptor(ctx, interceptorData{"data://", data}, handler)
}
dataServer := dataconn.NewServer(dataServerClientIdentitySetter, dataCtxInterceptor, loggers.Data, handler)
dataServerServe := func(ctx context.Context, dataListener transport.AuthenticatedListener, errOut chan<- error) {
dataServer.Serve(ctx, dataListener)
errOut <- nil // TODO bad design of dataServer?
}
server := &Server{
logger: loggers.General,
handler: handler,
controlServerServe: controlServerServe,
dataServer: dataServer,
dataServerServe: dataServerServe,
}
return server
}
// The context is used for cancellation only.
// Serve never returns an error, it logs them to the Server's logger.
func (s *Server) Serve(ctx context.Context, l transport.AuthenticatedListener) {
ctx, cancel := context.WithCancel(ctx)
defer cancel()
defer s.logger.Debug("rpc.(*Server).Serve done")
l = versionhandshake.Listener(l, envconst.Duration("ZREPL_RPC_SERVER_VERSIONHANDSHAKE_TIMEOUT", 10*time.Second))
// it is important that demux's context is cancelled,
// it has background goroutines attached
demuxListener := demux(ctx, l)
serveErrors := make(chan error, 2)
go s.controlServerServe(ctx, demuxListener.control, serveErrors)
go s.dataServerServe(ctx, demuxListener.data, serveErrors)
select {
case serveErr := <-serveErrors:
s.logger.WithError(serveErr).Error("serve error")
s.logger.Debug("wait for other server to shut down")
cancel()
secondServeErr := <-serveErrors
s.logger.WithError(secondServeErr).Error("serve error")
case <-ctx.Done():
s.logger.Debug("context cancelled, wait for control and data servers")
cancel()
for i := 0; i < 2; i++ {
<-serveErrors
}
s.logger.Debug("control and data server shut down, returning from Serve")
}
}
+254
View File
@@ -0,0 +1,254 @@
// Package transportmux wraps a transport.{Connecter,AuthenticatedListener}
// to distinguish different connection types based on a label
// sent from client to server on connection establishment.
//
// Labels are plain text and fixed length.
package transportmux
import (
"context"
"sync/atomic"
"syscall"
"fmt"
"io"
"net"
"time"
"github.com/zrepl/zrepl/internal/daemon/logging"
"github.com/zrepl/zrepl/internal/logger"
"github.com/zrepl/zrepl/internal/transport"
)
type Logger = logger.Logger
func getLog(ctx context.Context) Logger {
return logging.GetLogger(ctx, logging.SubsysTransportMux)
}
type acceptRes struct {
conn *transport.AuthConn
err error
}
type demuxListener struct {
closed int32
conns chan acceptRes
}
var ErrClosed = &net.OpError{
Op: "accept",
Net: "demux",
Source: nil,
Addr: nil,
Err: syscall.EINVAL,
}
func (l *demuxListener) Accept(ctx context.Context) (*transport.AuthConn, error) {
if atomic.LoadInt32(&l.closed) != 0 {
return nil, ErrClosed
}
select {
case r, ok := <-l.conns:
if !ok {
return nil, ErrClosed
}
return r.conn, r.err
case <-ctx.Done():
return nil, ctx.Err()
}
}
type demuxAddr struct{}
func (demuxAddr) Network() string { return "demux" }
func (demuxAddr) String() string { return "demux" }
func (l *demuxListener) Addr() net.Addr {
return demuxAddr{}
}
func (l *demuxListener) Close() error {
atomic.StoreInt32(&l.closed, 1)
return nil
}
// Exact length of a label in bytes (0-byte padded if it is shorter).
// This is a protocol constant, changing it breaks the wire protocol.
const LabelLen = 64
func padLabel(out []byte, label string) error {
if len(label) > LabelLen {
return fmt.Errorf("label %q exceeds max length (is %d, max %d)", label, len(label), LabelLen)
}
if len(out) != LabelLen {
panic(fmt.Sprintf("implementation error: %d", out))
}
labelBytes := []byte(label)
copy(out[:], labelBytes)
return nil
}
func Demux(ctx context.Context, rawListener transport.AuthenticatedListener, labels []string, timeout time.Duration) (map[string]transport.AuthenticatedListener, error) {
padded := make(map[[64]byte]*demuxListener, len(labels))
ret := make(map[string]transport.AuthenticatedListener, len(labels))
for _, label := range labels {
var labelPadded [LabelLen]byte
err := padLabel(labelPadded[:], label)
if err != nil {
return nil, err
}
if _, ok := padded[labelPadded]; ok {
return nil, fmt.Errorf("duplicate label %q", label)
}
dl := &demuxListener{
closed: 0,
conns: make(chan acceptRes, 1),
}
padded[labelPadded] = dl
ret[label] = dl
}
// invariant: padded contains same-length, non-duplicate labels
go func() {
<-ctx.Done()
getLog(ctx).Debug("context cancelled, closing listener")
if err := rawListener.Close(); err != nil {
getLog(ctx).WithError(err).Error("error closing listener")
}
drainConns := func(ch chan acceptRes) {
for c := range ch {
if c.conn != nil {
if err := c.conn.Close(); err != nil {
getLog(ctx).WithError(err).Error("error closing connection while draining after listener was closed")
}
}
}
}
for _, dl := range ret {
atomic.StoreInt32(&dl.(*demuxListener).closed, 1)
drainConns(dl.(*demuxListener).conns)
}
}()
go func() {
defer func() {
for _, dl := range ret {
close(dl.(*demuxListener).conns)
}
}()
for {
select {
case <-ctx.Done():
getLog(ctx).WithError(ctx.Err()).Info("stop accepting new connections after context done")
return
default:
}
rawConn, err := rawListener.Accept(ctx)
if err != nil {
if ctx.Err() != nil {
return
}
getLog(ctx).WithError(err).WithField("errType", fmt.Sprintf("%T", err)).Error("accept error")
continue
}
closeConn := func() {
if err := rawConn.Close(); err != nil {
getLog(ctx).WithError(err).Error("cannot close conn")
}
}
if err := rawConn.SetDeadline(time.Now().Add(timeout)); err != nil {
getLog(ctx).WithError(err).Error("SetDeadline failed")
closeConn()
continue
}
var labelBuf [LabelLen]byte
if _, err := io.ReadFull(rawConn, labelBuf[:]); err != nil {
getLog(ctx).WithError(err).Error("error reading label")
closeConn()
continue
}
demuxListener, ok := padded[labelBuf]
if !ok {
getLog(ctx).WithError(err).
WithField("client_label", fmt.Sprintf("%q", labelBuf)).
Error("unknown client label")
closeConn()
continue
}
err = rawConn.SetDeadline(time.Time{})
if err != nil {
getLog(ctx).WithError(err).Error("cannot reset deadline")
}
demuxListener.conns <- acceptRes{conn: rawConn, err: nil}
}
}()
return ret, nil
}
type labeledConnecter struct {
label []byte
transport.Connecter
}
func (c labeledConnecter) Connect(ctx context.Context) (transport.Wire, error) {
conn, err := c.Connecter.Connect(ctx)
if err != nil {
return nil, err
}
closeConn := func(why error) {
getLog(ctx).WithField("reason", why.Error()).Debug("closing connection")
if err := conn.Close(); err != nil {
getLog(ctx).WithError(err).Error("error closing connection after label write error")
}
}
if dl, ok := ctx.Deadline(); ok {
defer func() {
err := conn.SetDeadline(time.Time{})
if err != nil {
getLog(ctx).WithError(err).Error("cannot reset deadline")
}
}()
if err := conn.SetDeadline(dl); err != nil {
closeConn(err)
return nil, err
}
}
n, err := conn.Write(c.label)
if err != nil {
closeConn(err)
return nil, err
}
if n != len(c.label) {
closeConn(fmt.Errorf("short label write"))
return nil, io.ErrShortWrite
}
return conn, nil
}
func MuxConnecter(rawConnecter transport.Connecter, labels []string, timeout time.Duration) (map[string]transport.Connecter, error) {
ret := make(map[string]transport.Connecter, len(labels))
for _, label := range labels {
var paddedLabel [LabelLen]byte
if err := padLabel(paddedLabel[:], label); err != nil {
return nil, err
}
lc := &labeledConnecter{paddedLabel[:], rawConnecter}
if _, ok := ret[label]; ok {
return nil, fmt.Errorf("duplicate label %q", label)
}
ret[label] = lc
}
return ret, nil
}
@@ -0,0 +1,239 @@
// Package versionhandshake wraps a transport.{Connecter,AuthenticatedListener}
// to add an exchange of protocol version information on connection establishment.
//
// The protocol version information (banner) is plain text, thus making it
// easy to diagnose issues with standard tools.
package versionhandshake
import (
"bytes"
"fmt"
"io"
"net"
"strings"
"time"
"unicode/utf8"
)
type HandshakeMessage struct {
ProtocolVersion int
Extensions []string
}
// A HandshakeError describes what went wrong during the handshake.
// It implements net.Error and is always temporary.
type HandshakeError struct {
msg string
// If not nil, the underlying IO error that caused the handshake to fail.
IOError error
isAcceptError bool
}
var _ net.Error = &HandshakeError{}
func (e HandshakeError) Error() string { return e.msg }
// When a net.Listener.Accept() returns an error, the server must
// decide whether to retry calling Accept() or not.
// On some platforms (e.g., Linux), Accept() can return errors
// related to the specific protocol connection that was supposed
// to be returned as asocket FD. Obviously, we want to ignore,
// maybe log, those errors and retry Accept() immediately to
// serve other connections.
// But there are also conditions where we get Accept() errors because
// the process has run out of file descriptors. In that case, retrying
// won't help. We need to close some file descriptor to make progress.
// Note that there could be lots of open file descriptors because we
// have accepted, and not yet closed, lots of connections in the past.
// And then, of course there can be errors where we just want
// to return, e.g., if there's a programming error and we're getting
// an EBADFD or whatever.
//
// So, the serve loops in net/http.Server.Serve() or gRPC's server.Serve()
// must inspect the error and decide what to do.
// The vehicle for this is the
//
// interface { Temporary() bool }
//
// Behavior in both of the aforementioned Serve() loops:
//
// - if the error doesn't implement the interface, stop serving and return
// - `Temporary() == true`: retry with back-off
// - `Temporary() == false`: stop serving and return
//
// So, to make this package's HandshakeListener work with these
// Serve() loops, we return Temporary() == true if the handshake fails.
// In the aforementioned categories, that's the case of a per-connection
// protocol error.
//
// Note: the net.Error interface has deprecated the Temporary() method
// in go.dev/issue/45729, but there is no replacement for users of .Accept().
// Existing users of .Accept() continue to check for the interface.
// So, we need to continue supporting Temporary() until there's a different
// mechanism for serve loops to decide whether to retry or not.
// The following mailing list post proposes to eliminate the retries
// completely, but it seems like the effort has stalled.
// https://groups.google.com/g/golang-nuts/c/-JcZzOkyqYI/m/xwaZzjCgAwAJ
func (e HandshakeError) Temporary() bool {
if e.isAcceptError {
return true
}
te, ok := e.IOError.(interface{ Temporary() bool })
return ok && te.Temporary()
}
// If the underlying IOError was net.Error.Timeout(), Timeout() returns that value.
// Otherwise false.
func (e HandshakeError) Timeout() bool {
if neterr, ok := e.IOError.(net.Error); ok {
return neterr.Timeout()
}
return false
}
func hsErr(msg string) *HandshakeError {
return &HandshakeError{msg: msg}
}
func hsIOErr(err error, format string, args ...interface{}) *HandshakeError {
return &HandshakeError{IOError: err, msg: fmt.Sprintf(format, args...)}
}
// MaxProtocolVersion is the maximum allowed protocol version.
// This is a protocol constant, changing it may break the wire format.
const MaxProtocolVersion = 9999
// Only returns *HandshakeError as error.
func (m *HandshakeMessage) Encode() ([]byte, error) {
if m.ProtocolVersion <= 0 || m.ProtocolVersion > MaxProtocolVersion {
return nil, hsErr(fmt.Sprintf("protocol version must be in [1, %d]", MaxProtocolVersion))
}
if len(m.Extensions) >= MaxProtocolVersion {
return nil, hsErr(fmt.Sprintf("protocol only supports [0, %d] extensions", MaxProtocolVersion))
}
// EXTENSIONS is a count of subsequent \n separated lines that contain protocol extensions
var extensions strings.Builder
for i, ext := range m.Extensions {
if strings.ContainsAny(ext, "\n") {
return nil, hsErr(fmt.Sprintf("Extension #%d contains forbidden newline character", i))
}
if !utf8.ValidString(ext) {
return nil, hsErr(fmt.Sprintf("Extension #%d is not valid UTF-8", i))
}
extensions.WriteString(ext)
extensions.WriteString("\n")
}
withoutLen := fmt.Sprintf("ZREPL_ZFS_REPLICATION PROTOVERSION=%04d EXTENSIONS=%04d\n%s",
m.ProtocolVersion, len(m.Extensions), extensions.String())
withLen := fmt.Sprintf("%010d %s", len(withoutLen), withoutLen)
return []byte(withLen), nil
}
func (m *HandshakeMessage) DecodeReader(r io.Reader, maxLen int) error {
var lenAndSpace [11]byte
if _, err := io.ReadFull(r, lenAndSpace[:]); err != nil {
return hsIOErr(err, "error reading protocol banner length: %s", err)
}
if !utf8.Valid(lenAndSpace[:]) {
return hsErr("invalid start of handshake message: not valid UTF-8")
}
var followLen int
n, err := fmt.Sscanf(string(lenAndSpace[:]), "%010d ", &followLen)
if n != 1 || err != nil {
return hsErr("could not parse handshake message length")
}
if followLen > maxLen {
return hsErr(fmt.Sprintf("handshake message length exceeds max length (%d vs %d)",
followLen, maxLen))
}
var buf bytes.Buffer
_, err = io.Copy(&buf, io.LimitReader(r, int64(followLen)))
if err != nil {
return hsIOErr(err, "error reading protocol banner body: %s", err)
}
var (
protoVersion, extensionCount int
)
n, err = fmt.Fscanf(&buf, "ZREPL_ZFS_REPLICATION PROTOVERSION=%04d EXTENSIONS=%4d\n",
&protoVersion, &extensionCount)
if n != 2 || err != nil {
return hsErr(fmt.Sprintf("could not parse handshake message: %s", err))
}
if protoVersion < 1 {
return hsErr(fmt.Sprintf("invalid protocol version %q", protoVersion))
}
m.ProtocolVersion = protoVersion
if extensionCount < 0 {
return hsErr(fmt.Sprintf("invalid extension count %q", extensionCount))
}
if extensionCount == 0 {
if buf.Len() != 0 {
return hsErr("unexpected data trailing after header")
}
m.Extensions = nil
return nil
}
s := buf.String()
if strings.Count(s, "\n") != extensionCount {
return hsErr(fmt.Sprintf("inconsistent extension count: found %d, header says %d", len(m.Extensions), extensionCount))
}
exts := strings.Split(s, "\n")
if exts[len(exts)-1] != "" {
return hsErr("unexpected data trailing after last extension newline")
}
m.Extensions = exts[0 : len(exts)-1]
return nil
}
func DoHandshakeCurrentVersion(conn net.Conn, deadline time.Time) *HandshakeError {
// current protocol version is hardcoded here
return DoHandshakeVersion(conn, deadline, 7)
}
const HandshakeMessageMaxLen = 16 * 4096
func DoHandshakeVersion(conn net.Conn, deadline time.Time, version int) (rErr *HandshakeError) {
ours := HandshakeMessage{
ProtocolVersion: version,
Extensions: nil,
}
hsb, err := ours.Encode()
if err != nil {
return hsErr(fmt.Sprintf("could not encode protocol banner: %s", err))
}
err = conn.SetDeadline(deadline)
if err != nil {
return hsErr(fmt.Sprintf("could not set deadline for protocol banner handshake: %s", err))
}
defer func() {
if rErr != nil {
return
}
err := conn.SetDeadline(time.Time{})
if err != nil {
rErr = hsErr(fmt.Sprintf("could not reset deadline after protocol banner handshake: %s", err))
}
}()
_, err = io.Copy(conn, bytes.NewBuffer(hsb))
if err != nil {
return hsErr(fmt.Sprintf("could not send protocol banner: %s", err))
}
theirs := HandshakeMessage{}
if err := theirs.DecodeReader(conn, HandshakeMessageMaxLen); err != nil {
return hsErr(fmt.Sprintf("could not decode protocol banner: %s", err))
}
if theirs.ProtocolVersion != ours.ProtocolVersion {
return hsErr(fmt.Sprintf("protocol versions do not match: ours is %d, theirs is %d",
ours.ProtocolVersion, theirs.ProtocolVersion))
}
// ignore extensions, we don't use them
return nil
}
@@ -0,0 +1,119 @@
package versionhandshake
import (
"bytes"
"fmt"
"io"
"strings"
"testing"
"time"
"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/zrepl/zrepl/internal/util/socketpair"
)
func TestHandshakeMessage_Encode(t *testing.T) {
msg := HandshakeMessage{
ProtocolVersion: 2342,
}
encB, err := msg.Encode()
require.NoError(t, err)
enc := string(encB)
t.Logf("enc: %s", enc)
assert.False(t, strings.ContainsAny(enc[0:10], " "))
assert.True(t, enc[10] == ' ')
var (
headerlen, protoversion, extensionCount int
)
n, err := fmt.Sscanf(enc, "%010d ZREPL_ZFS_REPLICATION PROTOVERSION=%04d EXTENSIONS=%04d\n",
&headerlen, &protoversion, &extensionCount)
if n != 3 || (err != nil && err != io.EOF) {
t.Fatalf("%v %v", n, err)
}
assert.Equal(t, 2342, protoversion)
assert.Equal(t, 0, extensionCount)
assert.Equal(t, len(enc)-11, headerlen)
}
func TestHandshakeMessage_Encode_InvalidProtocolVersion(t *testing.T) {
for _, pv := range []int{-1, 0, 10000, 10001} {
t.Logf("testing invalid protocol version = %v", pv)
msg := HandshakeMessage{
ProtocolVersion: pv,
}
b, err := msg.Encode()
assert.Error(t, err)
assert.Nil(t, b)
}
}
func TestHandshakeMessage_DecodeReader(t *testing.T) {
in := HandshakeMessage{
2342,
[]string{"foo", "bar 2342"},
}
enc, err := in.Encode()
require.NoError(t, err)
out := HandshakeMessage{}
err = out.DecodeReader(bytes.NewReader([]byte(enc)), 4*4096)
assert.NoError(t, err)
assert.Equal(t, 2342, out.ProtocolVersion)
assert.Equal(t, 2, len(out.Extensions))
assert.Equal(t, "foo", out.Extensions[0])
assert.Equal(t, "bar 2342", out.Extensions[1])
}
func TestDoHandshakeVersion_ErrorOnDifferentVersions(t *testing.T) {
srv, client, err := socketpair.SocketPair()
if err != nil {
t.Fatal(err)
}
defer srv.Close()
defer client.Close()
srvErrCh := make(chan error)
go func() {
srvErrCh <- DoHandshakeVersion(srv, time.Now().Add(2*time.Second), 1)
}()
err = DoHandshakeVersion(client, time.Now().Add(2*time.Second), 2)
t.Log(err)
assert.Error(t, err)
assert.True(t, strings.Contains(err.Error(), "version"))
srvErr := <-srvErrCh
t.Log(srvErr)
assert.Error(t, srvErr)
assert.True(t, strings.Contains(srvErr.Error(), "version"))
}
func TestDoHandshakeCurrentVersion(t *testing.T) {
srv, client, err := socketpair.SocketPair()
if err != nil {
t.Fatal(err)
}
defer srv.Close()
defer client.Close()
srvErrCh := make(chan error)
go func() {
srvErrCh <- DoHandshakeVersion(srv, time.Now().Add(2*time.Second), 1)
}()
err = DoHandshakeVersion(client, time.Now().Add(2*time.Second), 1)
assert.Nil(t, err)
assert.Nil(t, <-srvErrCh)
}
@@ -0,0 +1,68 @@
package versionhandshake
import (
"context"
"net"
"time"
"github.com/zrepl/zrepl/internal/transport"
)
type HandshakeConnecter struct {
connecter transport.Connecter
timeout time.Duration
}
func (c HandshakeConnecter) Connect(ctx context.Context) (transport.Wire, error) {
conn, err := c.connecter.Connect(ctx)
if err != nil {
return nil, err
}
dl, ok := ctx.Deadline()
if !ok {
dl = time.Now().Add(c.timeout)
}
if err := DoHandshakeCurrentVersion(conn, dl); err != nil {
conn.Close()
return nil, err
}
return conn, nil
}
func Connecter(connecter transport.Connecter, timeout time.Duration) HandshakeConnecter {
return HandshakeConnecter{
connecter: connecter,
timeout: timeout,
}
}
// wrapper type that performs a a protocol version handshake before returning the connection
type HandshakeListener struct {
l transport.AuthenticatedListener
timeout time.Duration
}
func (l HandshakeListener) Addr() net.Addr { return l.l.Addr() }
func (l HandshakeListener) Close() error { return l.l.Close() }
func (l HandshakeListener) Accept(ctx context.Context) (*transport.AuthConn, error) {
conn, err := l.l.Accept(ctx)
if err != nil {
return nil, err
}
dl, ok := ctx.Deadline()
if !ok {
dl = time.Now().Add(l.timeout) // shadowing
}
if err := DoHandshakeCurrentVersion(conn, dl); err != nil {
err.isAcceptError = true
conn.Close()
return nil, err
}
return conn, nil
}
func Listener(l transport.AuthenticatedListener, timeout time.Duration) transport.AuthenticatedListener {
return HandshakeListener{l, timeout}
}