mirror of
https://github.com/fatedier/frp.git
synced 2026-10-15 15:55:58 +08:00
fix(server): control work connection refill scheduling
This commit is contained in:
fatedier
@@ -0,0 +1,739 @@
|
||||
// Copyright 2026 The frp Authors
|
||||
//
|
||||
// Licensed under the Apache License, Version 2.0 (the "License");
|
||||
// you may not use this file except in compliance with the License.
|
||||
// You may obtain a copy of the License at
|
||||
//
|
||||
// http://www.apache.org/licenses/LICENSE-2.0
|
||||
//
|
||||
// Unless required by applicable law or agreed to in writing, software
|
||||
// distributed under the License is distributed on an "AS IS" BASIS,
|
||||
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
// See the License for the specific language governing permissions and
|
||||
// limitations under the License.
|
||||
|
||||
package server
|
||||
|
||||
import (
|
||||
"container/list"
|
||||
"errors"
|
||||
"fmt"
|
||||
"runtime/debug"
|
||||
"slices"
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
pkgerr "github.com/fatedier/frp/pkg/errors"
|
||||
"github.com/fatedier/frp/pkg/msg"
|
||||
"github.com/fatedier/frp/pkg/util/xlog"
|
||||
"github.com/fatedier/frp/server/proxy"
|
||||
)
|
||||
|
||||
const workConnPoolCapacityOffset = 10
|
||||
|
||||
const (
|
||||
workConnReserveWindowMax = 8
|
||||
workConnDemandWindow = 64
|
||||
workConnLegacyWindow = 64
|
||||
|
||||
// Canceled waiters remain as lazy tombstones so cancellation stays O(1).
|
||||
// Compact only after enough tombstones have accumulated and comprise at
|
||||
// least half the queue, making the occasional O(n) copy O(1) amortized.
|
||||
workConnWaiterCompactMinTombstones = 64
|
||||
)
|
||||
|
||||
var (
|
||||
errWorkConnPoolFull = errors.New("work connection pool is full, discarding")
|
||||
errWorkConnControlUnavailable = errors.New("work connection control unavailable")
|
||||
)
|
||||
|
||||
type workConnRequestKind string
|
||||
|
||||
const (
|
||||
workConnRequestDemand workConnRequestKind = msg.WorkConnTypeDemand
|
||||
workConnRequestReserve workConnRequestKind = msg.WorkConnTypeReserve
|
||||
workConnRequestLegacy workConnRequestKind = "legacy"
|
||||
)
|
||||
|
||||
type workConnWaiter struct {
|
||||
resultCh chan workConnResult
|
||||
active bool
|
||||
handoffConn *proxy.WorkConn
|
||||
handoffClosed bool
|
||||
handoffElement *list.Element
|
||||
}
|
||||
|
||||
type workConnResult struct {
|
||||
conn *proxy.WorkConn
|
||||
err error
|
||||
}
|
||||
|
||||
type workConnLease struct {
|
||||
kind workConnRequestKind
|
||||
timer *time.Timer
|
||||
|
||||
active bool
|
||||
dispatched bool
|
||||
late bool
|
||||
}
|
||||
|
||||
// workConnScheduler owns all work-connection supply state. It deliberately
|
||||
// communicates with Control through callbacks for request dispatch and
|
||||
// dispatcher shutdown; it never acquires Control lifecycle or manager locks.
|
||||
type workConnScheduler struct {
|
||||
workConnCh chan *proxy.WorkConn
|
||||
|
||||
workConnMu sync.Mutex
|
||||
workConnWaiters []*workConnWaiter
|
||||
workConnActiveWaiters int
|
||||
workConnWaiterTombstones int
|
||||
// Elements retain registration order for cleanup; each waiter owns its
|
||||
// element so consumption removes it in O(1), even during large bursts.
|
||||
workConnHandoffs list.List
|
||||
workConnPending map[workConnRequestKind][]*workConnLease
|
||||
workConnLate map[workConnRequestKind][]*workConnLease
|
||||
workConnPoolClosed bool
|
||||
|
||||
workConnReconcileCh chan struct{}
|
||||
workConnReconcileStop chan struct{}
|
||||
workConnReconcileDone chan struct{}
|
||||
workConnReconcileOnce sync.Once
|
||||
workConnReconcileStopOnce sync.Once
|
||||
|
||||
workConnCleanupMu sync.Mutex
|
||||
workConnCleanupPending []*proxy.WorkConn
|
||||
workConnCleanupReady bool
|
||||
workConnCleanupStarted bool
|
||||
workConnCleanupDone chan struct{}
|
||||
|
||||
workConnLegacyMode bool
|
||||
workConnReserveWindow int
|
||||
workConnLeaseTimeout time.Duration
|
||||
poolCount int
|
||||
|
||||
// sendRequest is called only after workConnMu has been released. The
|
||||
// callback captures the Control's immutable generation metadata and sends
|
||||
// the wire message without exposing lifecycle state to the scheduler.
|
||||
sendRequest func(workConnRequestKind) error
|
||||
dispatcherDone func() <-chan struct{}
|
||||
logger *xlog.Logger
|
||||
doneCh <-chan struct{}
|
||||
}
|
||||
|
||||
func newWorkConnScheduler(poolCount int, leaseTimeout time.Duration, logger *xlog.Logger, doneCh <-chan struct{}) workConnScheduler {
|
||||
return workConnScheduler{
|
||||
workConnCh: make(chan *proxy.WorkConn, poolCount+workConnPoolCapacityOffset),
|
||||
poolCount: poolCount,
|
||||
workConnPending: make(map[workConnRequestKind][]*workConnLease),
|
||||
workConnLate: make(map[workConnRequestKind][]*workConnLease),
|
||||
workConnReconcileCh: make(chan struct{}, 1),
|
||||
workConnReconcileStop: make(chan struct{}),
|
||||
workConnReconcileDone: make(chan struct{}),
|
||||
workConnCleanupDone: make(chan struct{}),
|
||||
workConnLegacyMode: true,
|
||||
workConnReserveWindow: min(poolCount, workConnReserveWindowMax),
|
||||
workConnLeaseTimeout: leaseTimeout,
|
||||
logger: logger,
|
||||
doneCh: doneCh,
|
||||
}
|
||||
}
|
||||
|
||||
func normalizeWorkConnRequestKind(kind workConnRequestKind) workConnRequestKind {
|
||||
if kind == workConnRequestDemand || kind == workConnRequestReserve {
|
||||
return kind
|
||||
}
|
||||
return workConnRequestLegacy
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) workConnLeaseTimeoutValue() time.Duration {
|
||||
if s.workConnLeaseTimeout > 0 {
|
||||
return s.workConnLeaseTimeout
|
||||
}
|
||||
return 10 * time.Second
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) startWorkConnReconciler() {
|
||||
s.workConnReconcileOnce.Do(func() {
|
||||
go func() {
|
||||
defer close(s.workConnReconcileDone)
|
||||
for {
|
||||
select {
|
||||
case <-s.workConnReconcileCh:
|
||||
s.reconcileWorkConnRequests()
|
||||
case <-s.workConnReconcileStop:
|
||||
return
|
||||
}
|
||||
}
|
||||
}()
|
||||
})
|
||||
}
|
||||
|
||||
// requestWorkConnReconcile coalesces refill requests into one bounded trigger.
|
||||
// The reconciler may block on dispatcher backpressure, but it owns no control
|
||||
// lifecycle or generation locks and is stopped with the control pool.
|
||||
func (s *workConnScheduler) requestWorkConnReconcile() {
|
||||
s.workConnMu.Lock()
|
||||
closed := s.workConnPoolClosed
|
||||
s.workConnMu.Unlock()
|
||||
if closed {
|
||||
return
|
||||
}
|
||||
|
||||
s.startWorkConnReconciler()
|
||||
select {
|
||||
case <-s.workConnReconcileStop:
|
||||
return
|
||||
case s.workConnReconcileCh <- struct{}{}:
|
||||
default:
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) activeWaiterCountLocked() int {
|
||||
return s.workConnActiveWaiters
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) appendWorkConnWaiterLocked(waiter *workConnWaiter) {
|
||||
s.workConnWaiters = append(s.workConnWaiters, waiter)
|
||||
if waiter.active {
|
||||
s.workConnActiveWaiters++
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) popWorkConnWaiterLocked() *workConnWaiter {
|
||||
for len(s.workConnWaiters) > 0 {
|
||||
waiter := s.workConnWaiters[0]
|
||||
s.workConnWaiters[0] = nil
|
||||
s.workConnWaiters = s.workConnWaiters[1:]
|
||||
if waiter.active {
|
||||
waiter.active = false
|
||||
s.workConnActiveWaiters--
|
||||
if len(s.workConnWaiters) == 0 {
|
||||
s.workConnWaiters = nil
|
||||
}
|
||||
return waiter
|
||||
}
|
||||
if s.workConnWaiterTombstones > 0 {
|
||||
s.workConnWaiterTombstones--
|
||||
}
|
||||
}
|
||||
s.workConnWaiters = nil
|
||||
s.workConnWaiterTombstones = 0
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) compactWorkConnWaitersLocked() {
|
||||
if s.workConnWaiterTombstones == 0 {
|
||||
return
|
||||
}
|
||||
activeWaiters := make([]*workConnWaiter, 0, s.workConnActiveWaiters)
|
||||
for _, waiter := range s.workConnWaiters {
|
||||
if waiter.active {
|
||||
activeWaiters = append(activeWaiters, waiter)
|
||||
}
|
||||
}
|
||||
s.workConnWaiters = activeWaiters
|
||||
s.workConnWaiterTombstones = 0
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) shouldCompactWorkConnWaitersLocked() bool {
|
||||
tombstones := s.workConnWaiterTombstones
|
||||
if tombstones < workConnWaiterCompactMinTombstones {
|
||||
return false
|
||||
}
|
||||
return tombstones*2 >= len(s.workConnWaiters)
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) removeWorkConnHandoffLocked(target *workConnWaiter) {
|
||||
if target.handoffElement != nil {
|
||||
s.workConnHandoffs.Remove(target.handoffElement)
|
||||
target.handoffElement = nil
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) removeLeaseLocked(leases map[workConnRequestKind][]*workConnLease, target *workConnLease) {
|
||||
items := leases[target.kind]
|
||||
for i, lease := range items {
|
||||
if lease == target {
|
||||
items = append(items[:i], items[i+1:]...)
|
||||
if len(items) == 0 {
|
||||
delete(leases, target.kind)
|
||||
} else {
|
||||
leases[target.kind] = items
|
||||
}
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) hasLeaseLocked(leases map[workConnRequestKind][]*workConnLease, target *workConnLease) bool {
|
||||
return slices.Contains(leases[target.kind], target)
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) popLeaseLocked(leases map[workConnRequestKind][]*workConnLease, kind workConnRequestKind) *workConnLease {
|
||||
items := leases[kind]
|
||||
for len(items) > 0 {
|
||||
lease := items[0]
|
||||
items = items[1:]
|
||||
if len(items) == 0 {
|
||||
delete(leases, kind)
|
||||
} else {
|
||||
leases[kind] = items
|
||||
}
|
||||
if lease.active {
|
||||
return lease
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) leaseCountLocked(leases map[workConnRequestKind][]*workConnLease, kind workConnRequestKind) int {
|
||||
count := 0
|
||||
for _, lease := range leases[kind] {
|
||||
if lease.active {
|
||||
count++
|
||||
}
|
||||
}
|
||||
return count
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) workConnSupplyGapLocked() bool {
|
||||
target := s.activeWaiterCountLocked() + s.poolCount
|
||||
pending := s.leaseCountLocked(s.workConnPending, workConnRequestDemand) +
|
||||
s.leaseCountLocked(s.workConnPending, workConnRequestReserve)
|
||||
return len(s.workConnCh)+pending < target
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) consumeWorkConnLeaseLocked(kind workConnRequestKind) *workConnLease {
|
||||
kind = normalizeWorkConnRequestKind(kind)
|
||||
s.workConnLegacyMode = kind == workConnRequestLegacy
|
||||
kinds := []workConnRequestKind{kind}
|
||||
if s.workConnLegacyMode {
|
||||
// Untyped responses satisfy pending demand before pending reserve.
|
||||
kinds = []workConnRequestKind{workConnRequestDemand, workConnRequestReserve}
|
||||
}
|
||||
// There is no request ID: same-kind replies are interchangeable. A timely
|
||||
// retry must retire pending supply before late leases, otherwise its reply
|
||||
// is discarded while the retry remains pending and the pool never recovers.
|
||||
// For legacy clients, all pending kinds take precedence over late leases.
|
||||
for _, leases := range []map[workConnRequestKind][]*workConnLease{s.workConnPending, s.workConnLate} {
|
||||
for _, candidate := range kinds {
|
||||
if lease := s.popLeaseLocked(leases, candidate); lease != nil {
|
||||
lease.active = false
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
return lease
|
||||
}
|
||||
}
|
||||
}
|
||||
return nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) expireLateWorkConnLease(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if !lease.active || !lease.late {
|
||||
s.workConnMu.Unlock()
|
||||
return
|
||||
}
|
||||
lease.active = false
|
||||
s.removeLeaseLocked(s.workConnLate, lease)
|
||||
s.workConnMu.Unlock()
|
||||
s.requestWorkConnReconcile()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) expireWorkConnLease(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if !lease.active || lease.late {
|
||||
s.workConnMu.Unlock()
|
||||
return
|
||||
}
|
||||
s.removeLeaseLocked(s.workConnPending, lease)
|
||||
lease.late = true
|
||||
s.workConnLate[lease.kind] = append(s.workConnLate[lease.kind], lease)
|
||||
lease.timer = time.AfterFunc(s.workConnLeaseTimeoutValue(), func() {
|
||||
s.expireLateWorkConnLease(lease)
|
||||
})
|
||||
s.workConnMu.Unlock()
|
||||
s.requestWorkConnReconcile()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) newWorkConnLeaseLocked(kind workConnRequestKind) *workConnLease {
|
||||
lease := &workConnLease{
|
||||
kind: kind,
|
||||
active: true,
|
||||
}
|
||||
s.workConnPending[kind] = append(s.workConnPending[kind], lease)
|
||||
return lease
|
||||
}
|
||||
|
||||
// startWorkConnLeaseTimer starts the response timeout only after the request
|
||||
// has been accepted by Dispatcher.Send. A response may race this call; the
|
||||
// active/pending check under workConnMu prevents installing a timer on a lease
|
||||
// that was already consumed, failed, or closed.
|
||||
// Send acknowledges queue admission, not WriteMsg completion. The dispatcher
|
||||
// has no per-message write acknowledgement, so this timeout includes queue time.
|
||||
func (s *workConnScheduler) startWorkConnLeaseTimer(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if lease.active && !lease.dispatched && !lease.late && lease.timer == nil && s.hasLeaseLocked(s.workConnPending, lease) {
|
||||
lease.dispatched = true
|
||||
lease.timer = time.AfterFunc(s.workConnLeaseTimeoutValue(), func() {
|
||||
s.expireWorkConnLease(lease)
|
||||
})
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) registerWorkConn(conn *proxy.WorkConn, kind workConnRequestKind) (bool, error) {
|
||||
accepted, err := s.registerWorkConnState(conn, kind)
|
||||
s.reconcileWorkConnRequests()
|
||||
return accepted, err
|
||||
}
|
||||
|
||||
// registerWorkConnState commits work connection ownership and lease accounting
|
||||
// without sending refill requests. Callers that hold lifecycle or generation
|
||||
// locks must release them before invoking reconcileWorkConnRequests.
|
||||
func (s *workConnScheduler) registerWorkConnState(conn *proxy.WorkConn, kind workConnRequestKind) (bool, error) {
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return false, pkgerr.ErrCtlClosed
|
||||
}
|
||||
s.consumeWorkConnLeaseLocked(kind)
|
||||
if waiter := s.popWorkConnWaiterLocked(); waiter != nil {
|
||||
// A reserve response may satisfy a demand waiter, but any other
|
||||
// pending lease is still an in-flight source of supply until it really
|
||||
// times out. Keep it accounted for instead of creating a replacement.
|
||||
waiter.handoffConn = conn
|
||||
waiter.handoffElement = s.workConnHandoffs.PushBack(waiter)
|
||||
waiter.resultCh <- workConnResult{conn: conn}
|
||||
s.workConnMu.Unlock()
|
||||
return true, nil
|
||||
}
|
||||
// Any response that did not satisfy a waiter must still prove a real supply
|
||||
// gap before entering idle. This covers unmatched responses, demand leases
|
||||
// whose waiter timed out, and late leases whose replacement is already
|
||||
// pending. A matching lease has been consumed above, so pending includes
|
||||
// only the remaining supply credits; replies cannot identify their own lease.
|
||||
if !s.workConnSupplyGapLocked() {
|
||||
s.workConnMu.Unlock()
|
||||
return false, errWorkConnPoolFull
|
||||
}
|
||||
|
||||
select {
|
||||
case s.workConnCh <- conn:
|
||||
s.workConnMu.Unlock()
|
||||
return true, nil
|
||||
default:
|
||||
s.workConnMu.Unlock()
|
||||
return false, errWorkConnPoolFull
|
||||
}
|
||||
}
|
||||
|
||||
// consumeWorkConnWaiterResult transfers a direct waiter handoff to its
|
||||
// consumer only while the control pool is still open. A close/replacement may
|
||||
// have already detached the connection for cleanup even though the buffered
|
||||
// result notification remains readable; in that case the caller must never
|
||||
// receive the stale connection.
|
||||
func (s *workConnScheduler) consumeWorkConnWaiterResult(waiter *workConnWaiter, result workConnResult) (*proxy.WorkConn, error) {
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed || waiter.handoffClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
conn := waiter.handoffConn
|
||||
if conn != nil {
|
||||
waiter.handoffConn = nil
|
||||
s.removeWorkConnHandoffLocked(waiter)
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
if result.err != nil {
|
||||
return nil, result.err
|
||||
}
|
||||
if conn == nil {
|
||||
conn = result.conn
|
||||
}
|
||||
return conn, nil
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) cancelWorkConnWaiter(waiter *workConnWaiter) bool {
|
||||
canceled := false
|
||||
s.workConnMu.Lock()
|
||||
if waiter.active {
|
||||
waiter.active = false
|
||||
s.workConnActiveWaiters--
|
||||
s.workConnWaiterTombstones++
|
||||
if s.shouldCompactWorkConnWaitersLocked() {
|
||||
s.compactWorkConnWaitersLocked()
|
||||
}
|
||||
canceled = true
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
if canceled {
|
||||
s.requestWorkConnReconcile()
|
||||
}
|
||||
return canceled
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) startWorkConnCleanup() {
|
||||
s.workConnCleanupMu.Lock()
|
||||
if s.workConnCleanupStarted || !s.workConnCleanupReady {
|
||||
s.workConnCleanupMu.Unlock()
|
||||
return
|
||||
}
|
||||
conns := s.workConnCleanupPending
|
||||
s.workConnCleanupPending = nil
|
||||
s.workConnCleanupStarted = true
|
||||
s.workConnCleanupMu.Unlock()
|
||||
|
||||
if len(conns) == 0 {
|
||||
close(s.workConnCleanupDone)
|
||||
return
|
||||
}
|
||||
go func() {
|
||||
for _, conn := range conns {
|
||||
_ = conn.Close()
|
||||
}
|
||||
close(s.workConnCleanupDone)
|
||||
}()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) waitWorkConnCleanup() {
|
||||
s.workConnCleanupMu.Lock()
|
||||
started := s.workConnCleanupStarted
|
||||
s.workConnCleanupMu.Unlock()
|
||||
if started {
|
||||
<-s.workConnCleanupDone
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) takeWorkConnPool() {
|
||||
var conns []*proxy.WorkConn
|
||||
// Serialize pool extraction with cleanup startup. A concurrent worker may
|
||||
// observe the pool as closed before this caller has queued its streams;
|
||||
// keeping the cleanup gate held makes that ordering explicit and prevents
|
||||
// a premature workConnCleanupDone close from losing the extracted conns.
|
||||
s.workConnCleanupMu.Lock()
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
s.workConnCleanupMu.Unlock()
|
||||
return
|
||||
}
|
||||
s.workConnPoolClosed = true
|
||||
s.workConnReconcileStopOnce.Do(func() { close(s.workConnReconcileStop) })
|
||||
for _, waiter := range s.workConnWaiters {
|
||||
if waiter.active {
|
||||
waiter.active = false
|
||||
s.workConnActiveWaiters--
|
||||
waiter.resultCh <- workConnResult{err: pkgerr.ErrCtlClosed}
|
||||
}
|
||||
}
|
||||
s.workConnWaiters = nil
|
||||
s.workConnActiveWaiters = 0
|
||||
s.workConnWaiterTombstones = 0
|
||||
for _, leases := range s.workConnPending {
|
||||
for _, lease := range leases {
|
||||
lease.active = false
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
for _, leases := range s.workConnLate {
|
||||
for _, lease := range leases {
|
||||
lease.active = false
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
}
|
||||
}
|
||||
for element := s.workConnHandoffs.Front(); element != nil; element = element.Next() {
|
||||
waiter := element.Value.(*workConnWaiter)
|
||||
if waiter.handoffConn != nil {
|
||||
conns = append(conns, waiter.handoffConn)
|
||||
waiter.handoffConn = nil
|
||||
}
|
||||
waiter.handoffClosed = true
|
||||
waiter.handoffElement = nil
|
||||
}
|
||||
s.workConnHandoffs.Init()
|
||||
for {
|
||||
select {
|
||||
case conn := <-s.workConnCh:
|
||||
conns = append(conns, conn)
|
||||
default:
|
||||
close(s.workConnCh)
|
||||
s.workConnMu.Unlock()
|
||||
s.workConnCleanupPending = append(s.workConnCleanupPending, conns...)
|
||||
s.workConnCleanupReady = true
|
||||
s.workConnCleanupMu.Unlock()
|
||||
return
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) closeWorkConnPool() {
|
||||
s.takeWorkConnPool()
|
||||
s.startWorkConnCleanup()
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) reconcileWorkConnRequests() {
|
||||
type request struct {
|
||||
lease *workConnLease
|
||||
kind workConnRequestKind
|
||||
}
|
||||
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return
|
||||
}
|
||||
var requests []request
|
||||
demandPending := s.leaseCountLocked(s.workConnPending, workConnRequestDemand)
|
||||
reservePending := s.leaseCountLocked(s.workConnPending, workConnRequestReserve)
|
||||
activeWaiters := s.activeWaiterCountLocked()
|
||||
demandNeed := activeWaiters - demandPending
|
||||
if demandNeed > 0 {
|
||||
budget := workConnDemandWindow - demandPending
|
||||
if s.workConnLegacyMode {
|
||||
legacyBudget := workConnLegacyWindow - demandPending - reservePending
|
||||
if budget > legacyBudget {
|
||||
budget = legacyBudget
|
||||
}
|
||||
}
|
||||
if budget < 0 {
|
||||
budget = 0
|
||||
}
|
||||
if demandNeed > budget {
|
||||
demandNeed = budget
|
||||
}
|
||||
for i := 0; i < demandNeed; i++ {
|
||||
requests = append(requests, request{lease: s.newWorkConnLeaseLocked(workConnRequestDemand), kind: workConnRequestDemand})
|
||||
}
|
||||
demandPending += demandNeed
|
||||
}
|
||||
|
||||
// Keep the invariant available + pending >= waiters + poolCount, while
|
||||
// limiting only reserve requests to the sliding window. Late leases do not
|
||||
// consume retry budget; their responses are admitted separately above only
|
||||
// when this projected supply still has a real gap.
|
||||
reserveNeed := max(activeWaiters+s.poolCount-len(s.workConnCh)-demandPending-reservePending, 0)
|
||||
reserveBudget := s.workConnReserveWindow - reservePending
|
||||
if s.workConnLegacyMode {
|
||||
legacyBudget := workConnLegacyWindow - demandPending - reservePending
|
||||
if reserveBudget > legacyBudget {
|
||||
reserveBudget = legacyBudget
|
||||
}
|
||||
}
|
||||
if reserveBudget < 0 {
|
||||
reserveBudget = 0
|
||||
}
|
||||
if reserveNeed > reserveBudget {
|
||||
reserveNeed = reserveBudget
|
||||
}
|
||||
for i := 0; i < reserveNeed; i++ {
|
||||
requests = append(requests, request{lease: s.newWorkConnLeaseLocked(workConnRequestReserve), kind: workConnRequestReserve})
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
|
||||
for _, req := range requests {
|
||||
if err := s.sendRequest(req.kind); err != nil {
|
||||
s.failWorkConnLease(req.lease)
|
||||
if s.dispatcherDone != nil {
|
||||
select {
|
||||
case <-s.dispatcherDone():
|
||||
// Wake waiters immediately when the control transport is already
|
||||
// closed; otherwise they would wait until UserConnTimeout.
|
||||
s.closeWorkConnPool()
|
||||
default:
|
||||
}
|
||||
}
|
||||
} else {
|
||||
s.startWorkConnLeaseTimer(req.lease)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (s *workConnScheduler) failWorkConnLease(lease *workConnLease) {
|
||||
s.workConnMu.Lock()
|
||||
if lease.active && !lease.late {
|
||||
lease.active = false
|
||||
s.removeLeaseLocked(s.workConnPending, lease)
|
||||
if lease.timer != nil {
|
||||
lease.timer.Stop()
|
||||
}
|
||||
}
|
||||
s.workConnMu.Unlock()
|
||||
}
|
||||
|
||||
// When frps get one user connection, we get one work connection from the pool and return it.
|
||||
// If no workConn available in the pool, send message to frpc to get one or more
|
||||
// and wait until it is available.
|
||||
// return an error if wait timeout
|
||||
func (s *workConnScheduler) GetWorkConn() (workConn *proxy.WorkConn, err error) {
|
||||
xl := s.logger
|
||||
started := time.Now()
|
||||
defer func() {
|
||||
if recovered := recover(); recovered != nil {
|
||||
xl.Errorf("panic error: %v", recovered)
|
||||
xl.Errorf(string(debug.Stack()))
|
||||
workConn = nil
|
||||
err = pkgerr.ErrCtlClosed
|
||||
}
|
||||
}()
|
||||
|
||||
s.workConnMu.Lock()
|
||||
if s.workConnPoolClosed {
|
||||
s.workConnMu.Unlock()
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
// Get an idle connection first.
|
||||
select {
|
||||
case workConn = <-s.workConnCh:
|
||||
s.workConnMu.Unlock()
|
||||
xl.Debugf("get work connection from pool")
|
||||
s.requestWorkConnReconcile()
|
||||
return workConn, nil
|
||||
default:
|
||||
}
|
||||
|
||||
waiter := &workConnWaiter{
|
||||
resultCh: make(chan workConnResult, 1),
|
||||
active: true,
|
||||
}
|
||||
s.appendWorkConnWaiterLocked(waiter)
|
||||
s.workConnMu.Unlock()
|
||||
|
||||
// Pool hits need no timer. On a miss, include time already spent acquiring
|
||||
// the pool lock and install the timer before requesting a refill.
|
||||
timer := time.NewTimer(s.workConnLeaseTimeoutValue() - time.Since(started))
|
||||
defer timer.Stop()
|
||||
s.requestWorkConnReconcile()
|
||||
select {
|
||||
case result := <-waiter.resultCh:
|
||||
return s.consumeWorkConnWaiterResult(waiter, result)
|
||||
case <-timer.C:
|
||||
if !s.cancelWorkConnWaiter(waiter) {
|
||||
select {
|
||||
case result := <-waiter.resultCh:
|
||||
return s.consumeWorkConnWaiterResult(waiter, result)
|
||||
default:
|
||||
}
|
||||
}
|
||||
s.workConnMu.Lock()
|
||||
closed := s.workConnPoolClosed
|
||||
s.workConnMu.Unlock()
|
||||
if closed {
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
err = fmt.Errorf("timeout trying to get work connection")
|
||||
xl.Warnf("%v", err)
|
||||
return nil, err
|
||||
case <-s.doneCh:
|
||||
// doneCh is terminal for this control generation. Do not consume a
|
||||
// buffered successful handoff here: takeWorkConnPool detaches that
|
||||
// connection for cleanup, and returning it would leak it past shutdown.
|
||||
s.cancelWorkConnWaiter(waiter)
|
||||
return nil, pkgerr.ErrCtlClosed
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user