feat(notify): 重构通信内核并补齐 stream/bulk/record/transfer 能力
- 引入 LogicalConn/TransportConn 分层，ClientConn 保留兼容适配层 - 新增 Stream、Bulk、RecordStream 三条数据面能力及对应控制路径 - 完成 transfer/file 传输内核与状态快照、诊断能力 - 补齐 reconnect、inbound dispatcher、modern psk 等基础模块 - 增加大规模回归、并发与基准测试覆盖 - 更新依赖库
This commit is contained in:
@@ -0,0 +1,288 @@
|
||||
package notify
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"io"
|
||||
"sync"
|
||||
)
|
||||
|
||||
type StreamCopyOptions struct {
|
||||
BufferSize int
|
||||
CloseWrite bool
|
||||
CloseStream bool
|
||||
CloseWriter bool
|
||||
}
|
||||
|
||||
type StreamBridgeOptions struct {
|
||||
BufferSize int
|
||||
ClosePeerOnEOF bool
|
||||
ResetOnCopyError bool
|
||||
}
|
||||
|
||||
type StreamOpenCopyOptions struct {
|
||||
Open StreamOpenOptions
|
||||
Copy StreamCopyOptions
|
||||
}
|
||||
|
||||
func CopyToStream(ctx context.Context, stream Stream, src io.Reader, opt StreamCopyOptions) (int64, error) {
|
||||
if stream == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
if src == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
bufSize := opt.BufferSize
|
||||
if bufSize <= 0 {
|
||||
bufSize = defaultFileChunkSize
|
||||
}
|
||||
buf := make([]byte, bufSize)
|
||||
reader := newContextReader(ctx, src)
|
||||
written, err := io.CopyBuffer(stream, reader, buf)
|
||||
if err == nil || err == io.EOF {
|
||||
if opt.CloseWrite {
|
||||
if closeErr := stream.CloseWrite(); closeErr != nil {
|
||||
return written, closeErr
|
||||
}
|
||||
} else if opt.CloseStream {
|
||||
if closeErr := stream.Close(); closeErr != nil {
|
||||
return written, closeErr
|
||||
}
|
||||
}
|
||||
}
|
||||
return written, normalizeStreamCopyError(err)
|
||||
}
|
||||
|
||||
func CopyFromStream(ctx context.Context, dst io.Writer, stream Stream, opt StreamCopyOptions) (int64, error) {
|
||||
if stream == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
if dst == nil {
|
||||
return 0, io.ErrClosedPipe
|
||||
}
|
||||
bufSize := opt.BufferSize
|
||||
if bufSize <= 0 {
|
||||
bufSize = defaultFileChunkSize
|
||||
}
|
||||
buf := make([]byte, bufSize)
|
||||
reader := newContextReader(ctx, stream)
|
||||
written, err := io.CopyBuffer(dst, reader, buf)
|
||||
if (err == nil || err == io.EOF) && opt.CloseWriter {
|
||||
if closer, ok := dst.(io.Closer); ok {
|
||||
if closeErr := closer.Close(); closeErr != nil {
|
||||
return written, closeErr
|
||||
}
|
||||
}
|
||||
}
|
||||
return written, normalizeStreamCopyError(err)
|
||||
}
|
||||
|
||||
type contextReader struct {
|
||||
ctx context.Context
|
||||
src io.Reader
|
||||
}
|
||||
|
||||
func newContextReader(ctx context.Context, src io.Reader) io.Reader {
|
||||
if ctx == nil || src == nil {
|
||||
return src
|
||||
}
|
||||
return &contextReader{ctx: ctx, src: src}
|
||||
}
|
||||
|
||||
func (r *contextReader) Read(p []byte) (int, error) {
|
||||
if r == nil || r.src == nil {
|
||||
return 0, io.EOF
|
||||
}
|
||||
select {
|
||||
case <-r.ctx.Done():
|
||||
return 0, r.ctx.Err()
|
||||
default:
|
||||
}
|
||||
return r.src.Read(p)
|
||||
}
|
||||
|
||||
func normalizeStreamCopyError(err error) error {
|
||||
if err == io.EOF {
|
||||
return nil
|
||||
}
|
||||
return err
|
||||
}
|
||||
|
||||
func BridgeStream(ctx context.Context, stream Stream, peer io.ReadWriteCloser, opt StreamBridgeOptions) error {
|
||||
if stream == nil || peer == nil {
|
||||
return io.ErrClosedPipe
|
||||
}
|
||||
if ctx == nil {
|
||||
ctx = context.Background()
|
||||
}
|
||||
bridgeCtx, cancel := context.WithCancel(ctx)
|
||||
defer cancel()
|
||||
|
||||
var wg sync.WaitGroup
|
||||
errCh := make(chan error, 2)
|
||||
var abortOnce sync.Once
|
||||
var primaryErr error
|
||||
|
||||
abortBridge := func(err error) {
|
||||
abortOnce.Do(func() {
|
||||
primaryErr = err
|
||||
cancel()
|
||||
_ = peer.Close()
|
||||
if err != nil && opt.ResetOnCopyError {
|
||||
_ = stream.Reset(err)
|
||||
return
|
||||
}
|
||||
_ = stream.Close()
|
||||
})
|
||||
}
|
||||
|
||||
watchDone := make(chan struct{})
|
||||
watchStopped := make(chan struct{})
|
||||
go func() {
|
||||
defer close(watchStopped)
|
||||
select {
|
||||
case <-ctx.Done():
|
||||
abortBridge(ctx.Err())
|
||||
case <-watchDone:
|
||||
}
|
||||
}()
|
||||
|
||||
runCopy := func(fn func() error) {
|
||||
wg.Add(1)
|
||||
go func() {
|
||||
defer wg.Done()
|
||||
err := fn()
|
||||
if err != nil {
|
||||
abortBridge(err)
|
||||
}
|
||||
errCh <- err
|
||||
}()
|
||||
}
|
||||
|
||||
runCopy(func() error {
|
||||
_, err := CopyToStream(bridgeCtx, stream, peer, StreamCopyOptions{
|
||||
BufferSize: opt.BufferSize,
|
||||
CloseWrite: true,
|
||||
})
|
||||
if err != nil {
|
||||
cancel()
|
||||
}
|
||||
return err
|
||||
})
|
||||
|
||||
runCopy(func() error {
|
||||
_, err := copyFromStreamToBridgePeer(bridgeCtx, peer, stream, opt)
|
||||
if err != nil {
|
||||
cancel()
|
||||
}
|
||||
return err
|
||||
})
|
||||
|
||||
wg.Wait()
|
||||
close(errCh)
|
||||
close(watchDone)
|
||||
<-watchStopped
|
||||
|
||||
if primaryErr != nil {
|
||||
if ctx.Err() != nil {
|
||||
return ctx.Err()
|
||||
}
|
||||
return primaryErr
|
||||
}
|
||||
|
||||
var result error
|
||||
for err := range errCh {
|
||||
if err == nil {
|
||||
continue
|
||||
}
|
||||
if errors.Is(err, context.Canceled) && ctx.Err() == nil {
|
||||
continue
|
||||
}
|
||||
if result == nil {
|
||||
result = err
|
||||
}
|
||||
}
|
||||
return result
|
||||
}
|
||||
|
||||
type streamBridgeCloseWriter interface {
|
||||
CloseWrite() error
|
||||
}
|
||||
|
||||
func copyFromStreamToBridgePeer(ctx context.Context, peer io.ReadWriteCloser, stream Stream, opt StreamBridgeOptions) (int64, error) {
|
||||
written, err := CopyFromStream(ctx, peer, stream, StreamCopyOptions{
|
||||
BufferSize: opt.BufferSize,
|
||||
})
|
||||
if err != nil {
|
||||
return written, err
|
||||
}
|
||||
if closeWriter, ok := peer.(streamBridgeCloseWriter); ok {
|
||||
return written, closeWriter.CloseWrite()
|
||||
}
|
||||
if opt.ClosePeerOnEOF {
|
||||
return written, peer.Close()
|
||||
}
|
||||
return written, nil
|
||||
}
|
||||
|
||||
func OpenClientStreamFromReader(ctx context.Context, c Client, src io.Reader, opt StreamOpenCopyOptions) (Stream, int64, error) {
|
||||
if c == nil {
|
||||
return nil, 0, errStreamClientNil
|
||||
}
|
||||
return openStreamFromReader(ctx, src, opt, c.OpenStream)
|
||||
}
|
||||
|
||||
func OpenServerLogicalStreamFromReader(ctx context.Context, s Server, logical *LogicalConn, src io.Reader, opt StreamOpenCopyOptions) (Stream, int64, error) {
|
||||
if s == nil {
|
||||
return nil, 0, errStreamServerNil
|
||||
}
|
||||
if logical == nil {
|
||||
return nil, 0, errStreamLogicalConnNil
|
||||
}
|
||||
return openStreamFromReader(ctx, src, opt, func(ctx context.Context, openOpt StreamOpenOptions) (Stream, error) {
|
||||
return s.OpenStreamLogical(ctx, logical, openOpt)
|
||||
})
|
||||
}
|
||||
|
||||
func OpenServerTransportStreamFromReader(ctx context.Context, s Server, transport *TransportConn, src io.Reader, opt StreamOpenCopyOptions) (Stream, int64, error) {
|
||||
if s == nil {
|
||||
return nil, 0, errStreamServerNil
|
||||
}
|
||||
if transport == nil {
|
||||
return nil, 0, errStreamTransportNil
|
||||
}
|
||||
return openStreamFromReader(ctx, src, opt, func(ctx context.Context, openOpt StreamOpenOptions) (Stream, error) {
|
||||
return s.OpenStreamTransport(ctx, transport, openOpt)
|
||||
})
|
||||
}
|
||||
|
||||
func CopyStreamToWriter(ctx context.Context, stream Stream, dst io.Writer, opt StreamCopyOptions) (int64, error) {
|
||||
return CopyFromStream(ctx, dst, stream, opt)
|
||||
}
|
||||
|
||||
func openStreamFromReader(ctx context.Context, src io.Reader, opt StreamOpenCopyOptions, openFn func(context.Context, StreamOpenOptions) (Stream, error)) (Stream, int64, error) {
|
||||
if src == nil {
|
||||
return nil, 0, io.ErrClosedPipe
|
||||
}
|
||||
if openFn == nil {
|
||||
return nil, 0, io.ErrClosedPipe
|
||||
}
|
||||
opt = normalizeStreamOpenCopyOptions(opt)
|
||||
stream, err := openFn(ctx, opt.Open)
|
||||
if err != nil {
|
||||
return nil, 0, err
|
||||
}
|
||||
written, err := CopyToStream(ctx, stream, src, opt.Copy)
|
||||
if err != nil {
|
||||
_ = stream.Reset(err)
|
||||
return stream, written, err
|
||||
}
|
||||
return stream, written, nil
|
||||
}
|
||||
|
||||
func normalizeStreamOpenCopyOptions(opt StreamOpenCopyOptions) StreamOpenCopyOptions {
|
||||
if !opt.Copy.CloseWrite && !opt.Copy.CloseStream {
|
||||
opt.Copy.CloseWrite = true
|
||||
}
|
||||
return opt
|
||||
}
|
||||
Reference in New Issue
Block a user