Files
DragonTCP/core/cmd/dragontcp-client/chunk.go
T
2026-08-16 15:22:03 -03:00

809 lines
20 KiB
Go

package main
import (
"crypto/rand"
"encoding/binary"
"fmt"
"io"
"net"
"sort"
"sync"
"sync/atomic"
"time"
"dragontcp/internal/protocol"
"dragontcp/internal/wire"
)
type chunkClientOptions struct {
startSize int
minSize int
maxSize int
adaptive bool
adaptSuccesses int
adaptLog bool
pollers int
reconnectEvery int
pollDelay time.Duration
txnTimeout time.Duration
tcpBuffer int
minPipeline int
maxPipeline int
}
type adaptiveSizer struct {
mu sync.Mutex
name string
current int
min int
max int
adaptive bool
adaptSuccesses int
successes int
good int
bad int
logChanges bool
}
func newAdaptiveSizer(name string, start int, opts chunkClientOptions) *adaptiveSizer {
if start < opts.minSize {
start = opts.minSize
}
if start > opts.maxSize {
start = opts.maxSize
}
return &adaptiveSizer{
name: name,
current: start,
min: opts.minSize,
max: opts.maxSize,
adaptive: opts.adaptive,
adaptSuccesses: func() int {
if opts.adaptSuccesses > 0 {
return opts.adaptSuccesses
}
return 64
}(),
logChanges: opts.adaptLog,
}
}
func (s *adaptiveSizer) Current() int {
s.mu.Lock()
defer s.mu.Unlock()
return s.current
}
func (s *adaptiveSizer) Success(attempted int) {
s.mu.Lock()
defer s.mu.Unlock()
if !s.adaptive || attempted != s.current || s.current >= s.max {
return
}
if attempted > s.good {
s.good = attempted
}
s.successes++
growAfter := s.adaptSuccesses
if s.bad > 0 && s.bad-s.good <= 64 {
growAfter *= 8
}
if s.successes < growAfter {
return
}
s.successes = 0
old := s.current
next := 0
if s.bad > old+1 {
next = old + (s.bad-old)/2
} else {
if s.bad > 0 {
s.bad = 0
}
step := old / 4
if step < 32 {
step = 32
}
next = old + step
}
if next > s.max {
next = s.max
}
if next <= old {
return
}
s.current = next
if s.logChanges {
fmt.Printf("adaptive %s chunk: %d -> %d after stable success\n", s.name, old, next)
}
}
func (s *adaptiveSizer) Failure(attempted int) (int, int) {
s.mu.Lock()
defer s.mu.Unlock()
old := s.current
if !s.adaptive || attempted != s.current {
return old, old
}
s.successes = 0
if s.bad == 0 || attempted < s.bad {
s.bad = attempted
}
next := attempted / 2
if s.good > 0 && s.good < attempted {
next = s.good
} else {
s.good = 0
}
if next < s.min {
next = s.min
}
if next >= attempted && attempted > s.min {
next = attempted - 1
}
if next < s.min {
next = s.min
}
s.current = next
if s.logChanges && old != next {
fmt.Printf("adaptive %s chunk: %d -> %d after transport failure\n", s.name, old, next)
}
return old, next
}
type physicalConn struct {
conn net.Conn
requests int
}
type requestLane struct {
mu sync.Mutex
serverAddr string
tcpBuffer int
reconnectEvery int
timeout time.Duration
pc *physicalConn
closed bool
}
func newRequestLane(serverAddr string, tcpBuffer, reconnectEvery int, timeout time.Duration) *requestLane {
return &requestLane{
serverAddr: serverAddr,
tcpBuffer: tcpBuffer,
reconnectEvery: reconnectEvery,
timeout: timeout,
}
}
func (l *requestLane) discardLocked() {
if l.pc != nil {
_ = l.pc.conn.Close()
l.pc = nil
}
}
func (l *requestLane) closeAfterLocked() {
if l.pc != nil && l.reconnectEvery > 0 && l.pc.requests >= l.reconnectEvery {
l.discardLocked()
}
}
func (l *requestLane) ensureLocked() error {
if l.closed {
return net.ErrClosed
}
if l.pc != nil {
if l.reconnectEvery <= 0 || l.pc.requests < l.reconnectEvery {
return nil
}
l.discardLocked()
}
d := net.Dialer{Timeout: 10 * time.Second, KeepAlive: 30 * time.Second}
conn, err := d.Dial("tcp", l.serverAddr)
if err != nil {
return err
}
protocol.TuneTCP(conn)
protocol.TuneTCPBuffer(conn, l.tcpBuffer)
l.pc = &physicalConn{conn: conn}
return nil
}
func (l *requestLane) Close() {
l.mu.Lock()
l.closed = true
l.discardLocked()
l.mu.Unlock()
}
func (l *requestLane) single(mode byte, sid wire.SessionID, seq uint64, payload []byte) (byte, []byte, error) {
l.mu.Lock()
defer l.mu.Unlock()
if err := l.ensureLocked(); err != nil {
return 0, nil, err
}
timeout := l.timeout
if timeout <= 0 {
timeout = 5 * time.Second
}
_ = l.pc.conn.SetDeadline(time.Now().Add(timeout))
if err := wire.WriteRequest(l.pc.conn, mode, sid, seq, payload); err != nil {
l.discardLocked()
return 0, nil, err
}
status, body, err := wire.ReadResponse(l.pc.conn)
if err != nil {
l.discardLocked()
return 0, nil, err
}
l.pc.requests++
_ = l.pc.conn.SetDeadline(time.Time{})
l.closeAfterLocked()
if status != wire.StatusError && len(body) > 0 {
body = wire.DecodeMaskedResponse(status, body, sid, mode, seq)
}
return status, body, nil
}
// download sends one compact request and consumes up to count response records.
// startOffset is also the response keystream sequence. Each DATA response advances
// it by exactly the returned byte count.
func (l *requestLane) download(sid wire.SessionID, startOffset, ackOffset uint64, maxChunk, count int) ([][]byte, byte, error) {
l.mu.Lock()
defer l.mu.Unlock()
if err := l.ensureLocked(); err != nil {
return nil, 0, err
}
timeout := l.timeout
if timeout <= 0 {
timeout = 5 * time.Second
}
_ = l.pc.conn.SetDeadline(time.Now().Add(timeout))
payload := make([]byte, 14)
binary.BigEndian.PutUint64(payload[0:8], ackOffset)
binary.BigEndian.PutUint32(payload[8:12], uint32(maxChunk))
binary.BigEndian.PutUint16(payload[12:14], uint16(count))
if err := wire.WriteRequest(l.pc.conn, wire.ModeDownload, sid, startOffset, payload); err != nil {
l.discardLocked()
return nil, 0, err
}
out := make([][]byte, 0, count)
offset := startOffset
lastStatus := wire.StatusOK
for i := 0; i < count; i++ {
status, body, err := wire.ReadResponse(l.pc.conn)
if err != nil {
l.discardLocked()
return out, lastStatus, err
}
lastStatus = status
switch status {
case wire.StatusData:
body = wire.DecodeMaskedResponse(status, body, sid, wire.ModeDownload, offset)
if len(body) == 0 {
l.discardLocked()
return out, status, fmt.Errorf("empty DATA response")
}
out = append(out, body)
offset += uint64(len(body))
case wire.StatusWait, wire.StatusEOF:
i = count // stop after this response
case wire.StatusError:
l.discardLocked()
return out, status, fmt.Errorf("%s", string(body))
default:
l.discardLocked()
return out, status, fmt.Errorf("unknown response status %d", status)
}
if status == wire.StatusWait || status == wire.StatusEOF {
break
}
}
l.pc.requests++
_ = l.pc.conn.SetDeadline(time.Time{})
l.closeAfterLocked()
return out, lastStatus, nil
}
type pathProfile struct {
upload int
download int
persistent bool
at time.Time
}
var profileState struct {
sync.Mutex
key string
p pathProfile
}
var probeSeq atomic.Uint64
func randomSessionID() (wire.SessionID, error) {
var sid wire.SessionID
_, err := rand.Read(sid[:])
return sid, err
}
func probePattern(n int) []byte {
out := make([]byte, n)
for i := range out {
out[i] = byte((i*31 + 17) & 0xff)
}
return out
}
func makeProbePayload(kind byte, value, total int, token string) []byte {
base := 11 + len(token)
if total < base {
total = base
}
out := make([]byte, total)
copy(out[:4], wire.ProbeMagic[:])
out[4] = kind
binary.BigEndian.PutUint16(out[5:7], uint16(len(token)))
binary.BigEndian.PutUint32(out[7:11], uint32(value))
copy(out[11:11+len(token)], token)
for i := base; i < len(out); i++ {
out[i] = byte((i*31 + 17) & 0xff)
}
return out
}
func probeOne(serverAddr, token string, opts chunkClientOptions, kind byte, candidate int) bool {
sid, err := randomSessionID()
if err != nil {
return false
}
timeout := opts.txnTimeout
if timeout <= 0 || timeout > 2500*time.Millisecond {
timeout = 2500 * time.Millisecond
}
lane := newRequestLane(serverAddr, opts.tcpBuffer, 1, timeout)
defer lane.Close()
seq := probeSeq.Add(1)
total := 0
value := candidate
if kind == wire.ProbeUpload {
total = candidate
}
payload := makeProbePayload(kind, value, total, token)
status, body, err := lane.single(wire.ModeProbe, sid, seq, payload)
if err != nil {
return false
}
if kind == wire.ProbeUpload {
return status == wire.StatusOK
}
if kind == wire.ProbeDownload {
if status != wire.StatusData || len(body) != candidate {
return false
}
want := probePattern(candidate)
for i := range body {
if body[i] != want[i] {
return false
}
}
return true
}
return status == wire.StatusOK
}
func probePersistent(serverAddr, token string, opts chunkClientOptions) bool {
sid, err := randomSessionID()
if err != nil {
return false
}
timeout := opts.txnTimeout
if timeout <= 0 || timeout > 2500*time.Millisecond {
timeout = 2500 * time.Millisecond
}
lane := newRequestLane(serverAddr, opts.tcpBuffer, 0, timeout)
defer lane.Close()
for i := 0; i < 8; i++ {
seq := probeSeq.Add(1)
payload := makeProbePayload(wire.ProbeKeepalive, i, 32+len(token), token)
status, _, err := lane.single(wire.ModeProbe, sid, seq, payload)
if err != nil || status != wire.StatusOK {
return false
}
}
return true
}
func probeCandidates(minSize, maxSize int) []int {
base := []int{32, 64, 128, 256, 512, 1024, 1200, 1280, 1320, 1350, 1360, 1380, 1400, 1450, 1600, 2048, 3205, 4096, 8192, 16384, 32768, 65536, 98304, 131072, 262144, 524288, 786432, 1048576}
seen := map[int]bool{}
out := make([]int, 0, len(base)+2)
for _, n := range base {
if n >= minSize && n <= maxSize && !seen[n] {
out = append(out, n)
seen[n] = true
}
}
if !seen[minSize] {
out = append(out, minSize)
}
if !seen[maxSize] {
out = append(out, maxSize)
}
sort.Ints(out)
return out
}
func probeMaximum(serverAddr, token string, opts chunkClientOptions, kind byte) int {
candidates := probeCandidates(opts.minSize, opts.maxSize)
lo, hi := 0, len(candidates)-1
best := opts.minSize
for lo <= hi {
mid := lo + (hi-lo)/2
candidate := candidates[mid]
if probeOne(serverAddr, token, opts, kind, candidate) {
best = candidate
lo = mid + 1
} else {
hi = mid - 1
}
}
if best < opts.minSize {
best = opts.minSize
}
return best
}
func getPathProfile(serverAddr, token string, opts chunkClientOptions) pathProfile {
key := fmt.Sprintf("%s|%s|%d|%d", serverAddr, token, opts.minSize, opts.maxSize)
profileState.Lock()
if profileState.key == key && time.Since(profileState.p.at) < 30*time.Minute {
p := profileState.p
profileState.Unlock()
return p
}
profileState.Unlock()
fallbackUp := minInt(opts.maxSize, maxInt(opts.minSize, 32768))
fallbackDown := minInt(opts.maxSize, maxInt(opts.minSize, 1350))
upCh := make(chan int, 1)
downCh := make(chan int, 1)
go func() { upCh <- probeMaximum(serverAddr, token, opts, wire.ProbeUpload) }()
go func() { downCh <- probeMaximum(serverAddr, token, opts, wire.ProbeDownload) }()
p := pathProfile{upload: fallbackUp, download: fallbackDown, persistent: false, at: time.Now()}
select {
case p.upload = <-upCh:
case <-time.After(20 * time.Second):
}
select {
case p.download = <-downCh:
case <-time.After(20 * time.Second):
}
p.persistent = probePersistent(serverAddr, token, opts)
fmt.Printf("path probe: upload=%d download=%d persistent=%t\n", p.upload, p.download, p.persistent)
profileState.Lock()
profileState.key = key
profileState.p = p
profileState.Unlock()
return p
}
func encodeOpen(token, host string, port int) ([]byte, error) {
if len(token) > 65535 || len(host) > 65535 {
return nil, fmt.Errorf("token or hostname too long")
}
out := make([]byte, 6+len(token)+len(host))
binary.BigEndian.PutUint16(out[0:2], uint16(len(token)))
binary.BigEndian.PutUint16(out[2:4], uint16(len(host)))
binary.BigEndian.PutUint16(out[4:6], uint16(port))
copy(out[6:6+len(token)], token)
copy(out[6+len(token):], host)
return out, nil
}
type chunkConn struct {
sid wire.SessionID
opts chunkClientOptions
uploadLane *requestLane
downloadLane *requestLane
serverMax int
upSizer *adaptiveSizer
downSizer *adaptiveSizer
writeMu sync.Mutex
upOffset uint64
readMu sync.Mutex
readBuf []byte
downloadOffset uint64
consumedOffset uint64
eof bool
pipeline int
minPipeline int
maxPipeline int
closeOnce sync.Once
}
func openChunkTunnel(serverAddr, token, targetHost string, targetPort int, opts chunkClientOptions) (net.Conn, error) {
if opts.minSize < 32 {
opts.minSize = 32
}
if opts.maxSize < opts.minSize {
opts.maxSize = opts.minSize
}
if opts.maxSize > 1024*1024 {
opts.maxSize = 1024 * 1024
}
if opts.txnTimeout <= 0 {
opts.txnTimeout = 5 * time.Second
}
if opts.adaptSuccesses < 1 {
opts.adaptSuccesses = 64
}
if opts.reconnectEvery < 0 {
opts.reconnectEvery = 0
}
if opts.maxPipeline < 1 {
opts.maxPipeline = 1
}
if opts.maxPipeline > 256 {
opts.maxPipeline = 256
}
if opts.minPipeline < 1 {
opts.minPipeline = 1
}
if opts.minPipeline > opts.maxPipeline {
opts.minPipeline = opts.maxPipeline
}
profile := getPathProfile(serverAddr, token, opts)
reconnect := opts.reconnectEvery
// Compatibility-friendly reconnect modes:
// 0 = persistent (CLI explicit)
// 1 = auto: persistent when the path probe succeeds, otherwise one request/connection
// N>=2 = force connection rotation after N logical requests
// Resolved silently: this runs once per proxied flow, so it must never log.
if reconnect == 1 && profile.persistent {
reconnect = 0
}
sid, err := randomSessionID()
if err != nil {
return nil, err
}
// OPEN rides the upload lane instead of a throwaway connection. A dedicated
// control connection cost one extra dial per proxied flow, which shows up on
// the server as connection churn on top of the steady-state count.
uploadLane := newRequestLane(serverAddr, opts.tcpBuffer, reconnect, opts.txnTimeout)
payload, err := encodeOpen(token, targetHost, targetPort)
if err != nil {
uploadLane.Close()
return nil, err
}
status, body, err := uploadLane.single(wire.ModeOpen, sid, 0, payload)
if err != nil {
uploadLane.Close()
return nil, err
}
if status == wire.StatusError {
uploadLane.Close()
return nil, fmt.Errorf("%s", string(body))
}
if status != wire.StatusOK || len(body) != 4 {
uploadLane.Close()
return nil, fmt.Errorf("bad OPEN response")
}
serverMax := int(binary.BigEndian.Uint32(body))
if serverMax < opts.minSize {
uploadLane.Close()
return nil, fmt.Errorf("server maximum chunk %d is below client minimum %d", serverMax, opts.minSize)
}
if opts.maxSize > serverMax {
opts.maxSize = serverMax
}
upStart := minInt(profile.upload, opts.maxSize)
downStart := minInt(profile.download, opts.maxSize)
if upStart < opts.minSize {
upStart = opts.minSize
}
if downStart < opts.minSize {
downStart = opts.minSize
}
c := &chunkConn{
sid: sid,
opts: opts,
serverMax: serverMax,
uploadLane: uploadLane,
downloadLane: newRequestLane(serverAddr, opts.tcpBuffer, reconnect, opts.txnTimeout),
// Start at the configured ceiling. On transport failure the batch is
// halved but never below minPipeline; successful data grows it back by
// one. When min == max the depth is pinned and never adapts, which is
// what paths that only work at one specific batch size need.
pipeline: opts.maxPipeline,
minPipeline: opts.minPipeline,
maxPipeline: opts.maxPipeline,
}
c.upSizer = newAdaptiveSizer("upload", upStart, opts)
c.downSizer = newAdaptiveSizer("download", downStart, opts)
return c, nil
}
func (c *chunkConn) fillReadBuffer() error {
if c.eof {
return io.EOF
}
minFailures := 0
for len(c.readBuf) == 0 && !c.eof {
chunk := c.downSizer.Current()
count := c.pipeline
if count < c.minPipeline {
count = c.minPipeline
}
if count > c.maxPipeline {
count = c.maxPipeline
}
// Bound each batch to roughly 1 MiB of useful data, but never below the
// configured floor: a pinned depth is a path requirement, not a hint.
if maxCount := (1024 * 1024) / maxInt(chunk, 1); maxCount < count {
count = maxInt(maxCount, c.minPipeline)
}
data, status, err := c.downloadLane.download(c.sid, c.downloadOffset, c.consumedOffset, chunk, count)
for _, part := range data {
c.readBuf = append(c.readBuf, part...)
c.downloadOffset += uint64(len(part))
}
if len(data) > 0 {
c.downSizer.Success(chunk)
if c.pipeline < c.maxPipeline {
c.pipeline++
}
minFailures = 0
}
if err != nil {
if c.pipeline > c.minPipeline {
old := c.pipeline
c.pipeline /= 2
if c.pipeline < c.minPipeline {
c.pipeline = c.minPipeline
}
if c.opts.adaptLog && old != c.pipeline {
fmt.Printf("adaptive download pipeline: %d -> %d after transport failure\n", old, c.pipeline)
}
} else {
old, next := c.downSizer.Failure(chunk)
if old == next && next == c.opts.minSize {
minFailures++
if minFailures >= 8 {
return fmt.Errorf("download failed at minimum chunk %d: %w", next, err)
}
}
}
time.Sleep(30 * time.Millisecond)
if len(c.readBuf) > 0 {
return nil
}
continue
}
switch status {
case wire.StatusEOF:
c.eof = true
case wire.StatusWait:
if c.opts.pollDelay > 0 {
time.Sleep(c.opts.pollDelay)
} else {
time.Sleep(5 * time.Millisecond)
}
}
if len(c.readBuf) > 0 {
return nil
}
}
if c.eof && len(c.readBuf) == 0 {
return io.EOF
}
return nil
}
func (c *chunkConn) Read(p []byte) (int, error) {
c.readMu.Lock()
defer c.readMu.Unlock()
if len(p) == 0 {
return 0, nil
}
if len(c.readBuf) == 0 {
if err := c.fillReadBuffer(); err != nil {
return 0, err
}
}
n := copy(p, c.readBuf)
c.readBuf = c.readBuf[n:]
c.consumedOffset += uint64(n)
return n, nil
}
func (c *chunkConn) Write(p []byte) (int, error) {
c.writeMu.Lock()
defer c.writeMu.Unlock()
if len(p) == 0 {
return 0, nil
}
total := 0
minFailures := 0
for len(p) > 0 {
size := c.upSizer.Current()
n := minInt(size, len(p))
status, body, err := c.uploadLane.single(wire.ModeUpload, c.sid, c.upOffset, p[:n])
if err != nil {
old, next := c.upSizer.Failure(size)
if old == next && next == c.opts.minSize {
minFailures++
if minFailures >= 8 {
return total, fmt.Errorf("upload failed at minimum chunk %d: %w", next, err)
}
} else {
minFailures = 0
}
time.Sleep(30 * time.Millisecond)
continue
}
if status == wire.StatusError {
return total, fmt.Errorf("%s", string(body))
}
if status != wire.StatusOK {
return total, fmt.Errorf("unexpected upload status %d", status)
}
c.upOffset += uint64(n)
total += n
p = p[n:]
c.upSizer.Success(size)
minFailures = 0
}
return total, nil
}
func (c *chunkConn) Close() error {
c.closeOnce.Do(func() {
// Reuse the upload lane rather than dialling a connection just to say
// goodbye; that was a second wasted dial per flow.
_, _, _ = c.uploadLane.single(wire.ModeClose, c.sid, 0, nil)
c.uploadLane.Close()
c.downloadLane.Close()
})
return nil
}
func (c *chunkConn) LocalAddr() net.Addr { return dummyAddr("dragontcp-binary-local") }
func (c *chunkConn) RemoteAddr() net.Addr { return dummyAddr("dragontcp-binary-remote") }
func (c *chunkConn) SetDeadline(time.Time) error { return nil }
func (c *chunkConn) SetReadDeadline(time.Time) error { return nil }
func (c *chunkConn) SetWriteDeadline(time.Time) error { return nil }
type dummyAddr string
func (d dummyAddr) Network() string { return "dragontcp-binary" }
func (d dummyAddr) String() string { return string(d) }
func minInt(a, b int) int {
if a < b {
return a
}
return b
}
func maxInt(a, b int) int {
if a > b {
return a
}
return b
}