chore: tuic use a simple client pool

This commit is contained in:
wwqgtxx
2022-11-25 20:14:05 +08:00
parent c7bad89af3
commit 7b44cde4bd
2 changed files with 102 additions and 57 deletions

View File

@ -39,13 +39,14 @@ type Client struct {
RequestTimeout int
MaxUdpRelayPacketSize int
Key any
LastVisited time.Time
UDP bool
quicConn quic.Connection
connMutex sync.Mutex
openStreams atomic.Int32
OpenStreams atomic.Int32
udpInputMap sync.Map
}
@ -242,9 +243,9 @@ func (t *Client) DialContext(ctx context.Context, metadata *C.Metadata, dialFn f
if err != nil {
return nil, err
}
openStreams := t.openStreams.Add(1)
openStreams := t.OpenStreams.Add(1)
if openStreams >= MaxOpenStreams {
t.openStreams.Add(-1)
t.OpenStreams.Add(-1)
return nil, TooManyOpenStreams
}
stream, err := func() (stream *quicStreamConn, err error) {
@ -300,6 +301,9 @@ type quicStreamConn struct {
lAddr net.Addr
rAddr net.Addr
client *Client
closeOnce sync.Once
closeErr error
}
func (q *quicStreamConn) Write(p []byte) (n int, err error) {
@ -309,8 +313,15 @@ func (q *quicStreamConn) Write(p []byte) (n int, err error) {
}
func (q *quicStreamConn) Close() error {
q.closeOnce.Do(func() {
q.closeErr = q.close()
})
return q.closeErr
}
func (q *quicStreamConn) close() error {
defer time.AfterFunc(C.DefaultTCPTimeout, func() {
q.client.openStreams.Add(-1)
q.client.OpenStreams.Add(-1)
})
// https://github.com/cloudflare/cloudflared/commit/ed2bac026db46b239699ac5ce4fcf122d7cab2cd
@ -342,6 +353,11 @@ func (t *Client) ListenPacketContext(ctx context.Context, metadata *C.Metadata,
if err != nil {
return nil, err
}
openStreams := t.OpenStreams.Add(1)
if openStreams >= MaxOpenStreams {
t.OpenStreams.Add(-1)
return nil, TooManyOpenStreams
}
pipe1, pipe2 := net.Pipe()
inputCh := make(chan udpData)
@ -380,16 +396,21 @@ type quicStreamPacketConn struct {
closeOnce sync.Once
closeErr error
closed bool
}
func (q *quicStreamPacketConn) Close() error {
q.closeOnce.Do(func() {
q.closed = true
q.closeErr = q.close()
})
return q.closeErr
}
func (q *quicStreamPacketConn) close() (err error) {
defer time.AfterFunc(C.DefaultTCPTimeout, func() {
q.client.OpenStreams.Add(-1)
})
defer func() {
q.client.deferQuicConn(q.quicConn, err)
}()
@ -441,6 +462,9 @@ func (q *quicStreamPacketConn) WriteTo(p []byte, addr net.Addr) (n int, err erro
if len(p) > q.client.MaxUdpRelayPacketSize {
return 0, fmt.Errorf("udp packet too large(%d > %d)", len(p), q.client.MaxUdpRelayPacketSize)
}
if q.closed {
return 0, net.ErrClosed
}
defer func() {
q.client.deferQuicConn(q.quicConn, err)
}()