Code refactoring related to message handling and retry logic. (#3745)

pkg/metrics/metrics.go

@@ -1,3 +1,17 @@
+// Copyright 2023 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 metrics
 
 import (

pkg/msg/handler.go

@@ -0,0 +1,103 @@
+// Copyright 2023 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 msg
+
+import (
+	"io"
+	"reflect"
+)
+
+func AsyncHandler(f func(Message)) func(Message) {
+	return func(m Message) {
+		go f(m)
+	}
+}
+
+// Dispatcher is used to send messages to net.Conn or register handlers for messages read from net.Conn.
+type Dispatcher struct {
+	rw io.ReadWriter
+
+	sendCh         chan Message
+	doneCh         chan struct{}
+	msgHandlers    map[reflect.Type]func(Message)
+	defaultHandler func(Message)
+}
+
+func NewDispatcher(rw io.ReadWriter) *Dispatcher {
+	return &Dispatcher{
+		rw:          rw,
+		sendCh:      make(chan Message, 100),
+		doneCh:      make(chan struct{}),
+		msgHandlers: make(map[reflect.Type]func(Message)),
+	}
+}
+
+// Run will block until io.EOF or some error occurs.
+func (d *Dispatcher) Run() {
+	go d.sendLoop()
+	go d.readLoop()
+}
+
+func (d *Dispatcher) sendLoop() {
+	for {
+		select {
+		case <-d.doneCh:
+			return
+		case m := <-d.sendCh:
+			_ = WriteMsg(d.rw, m)
+		}
+	}
+}
+
+func (d *Dispatcher) readLoop() {
+	for {
+		m, err := ReadMsg(d.rw)
+		if err != nil {
+			close(d.doneCh)
+			return
+		}
+
+		if handler, ok := d.msgHandlers[reflect.TypeOf(m)]; ok {
+			handler(m)
+		} else if d.defaultHandler != nil {
+			d.defaultHandler(m)
+		}
+	}
+}
+
+func (d *Dispatcher) Send(m Message) error {
+	select {
+	case <-d.doneCh:
+		return io.EOF
+	case d.sendCh <- m:
+		return nil
+	}
+}
+
+func (d *Dispatcher) SendChannel() chan Message {
+	return d.sendCh
+}
+
+func (d *Dispatcher) RegisterHandler(msg Message, handler func(Message)) {
+	d.msgHandlers[reflect.TypeOf(msg)] = handler
+}
+
+func (d *Dispatcher) RegisterDefaultHandler(handler func(Message)) {
+	d.defaultHandler = handler
+}
+
+func (d *Dispatcher) Done() chan struct{} {
+	return d.doneCh
+}

pkg/transport/message.go

@@ -29,7 +29,9 @@ type MessageTransporter interface {
 	// Recv(ctx context.Context, laneKey string, msgType string) (Message, error)
 	// Do will first send msg, then recv msg with the same laneKey and specified msgType.
 	Do(ctx context.Context, req msg.Message, laneKey, recvMsgType string) (msg.Message, error)
+	// Dispatch will dispatch message to releated channel registered in Do function by its message type and laneKey.
 	Dispatch(m msg.Message, laneKey string) bool
+	// Same with Dispatch but with specified message type.
 	DispatchWithType(m msg.Message, msgType, laneKey string) bool
 }
 

pkg/util/net/conn.go

@@ -22,6 +22,7 @@ import (
 	"sync/atomic"
 	"time"
 
+	"github.com/fatedier/golib/crypto"
 	quic "github.com/quic-go/quic-go"
 
 	"github.com/fatedier/frp/pkg/util/xlog"
@@ -216,3 +217,18 @@ func (conn *wrapQuicStream) Close() error {
 	conn.Stream.CancelRead(0)
 	return conn.Stream.Close()
 }
+
+func NewCryptoReadWriter(rw io.ReadWriter, key []byte) (io.ReadWriter, error) {
+	encReader := crypto.NewReader(rw, key)
+	encWriter, err := crypto.NewWriter(rw, key)
+	if err != nil {
+		return nil, err
+	}
+	return struct {
+		io.Reader
+		io.Writer
+	}{
+		Reader: encReader,
+		Writer: encWriter,
+	}, nil
+}

pkg/util/wait/backoff.go

@@ -0,0 +1,197 @@
+// Copyright 2023 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 wait
+
+import (
+	"math/rand"
+	"time"
+
+	"github.com/samber/lo"
+
+	"github.com/fatedier/frp/pkg/util/util"
+)
+
+type BackoffFunc func(previousDuration time.Duration, previousConditionError bool) time.Duration
+
+func (f BackoffFunc) Backoff(previousDuration time.Duration, previousConditionError bool) time.Duration {
+	return f(previousDuration, previousConditionError)
+}
+
+type BackoffManager interface {
+	Backoff(previousDuration time.Duration, previousConditionError bool) time.Duration
+}
+
+type FastBackoffOptions struct {
+	Duration           time.Duration
+	Factor             float64
+	Jitter             float64
+	MaxDuration        time.Duration
+	InitDurationIfFail time.Duration
+
+	// If FastRetryCount > 0, then within the FastRetryWindow time window,
+	// the retry will be performed with a delay of FastRetryDelay for the first FastRetryCount calls.
+	FastRetryCount  int
+	FastRetryDelay  time.Duration
+	FastRetryJitter float64
+	FastRetryWindow time.Duration
+}
+
+type fastBackoffImpl struct {
+	options FastBackoffOptions
+
+	lastCalledTime      time.Time
+	consecutiveErrCount int
+
+	fastRetryCutoffTime     time.Time
+	countsInFastRetryWindow int
+}
+
+func NewFastBackoffManager(options FastBackoffOptions) BackoffManager {
+	return &fastBackoffImpl{
+		options:                 options,
+		countsInFastRetryWindow: 1,
+	}
+}
+
+func (f *fastBackoffImpl) Backoff(previousDuration time.Duration, previousConditionError bool) time.Duration {
+	if f.lastCalledTime.IsZero() {
+		f.lastCalledTime = time.Now()
+		return f.options.Duration
+	}
+	now := time.Now()
+	f.lastCalledTime = now
+
+	if previousConditionError {
+		f.consecutiveErrCount++
+	} else {
+		f.consecutiveErrCount = 0
+	}
+
+	if f.options.FastRetryCount > 0 && previousConditionError {
+		f.countsInFastRetryWindow++
+		if f.countsInFastRetryWindow <= f.options.FastRetryCount {
+			return Jitter(f.options.FastRetryDelay, f.options.FastRetryJitter)
+		}
+		if now.After(f.fastRetryCutoffTime) {
+			// reset
+			f.fastRetryCutoffTime = now.Add(f.options.FastRetryWindow)
+			f.countsInFastRetryWindow = 0
+		}
+	}
+
+	if previousConditionError {
+		var duration time.Duration
+		if f.consecutiveErrCount == 1 {
+			duration = util.EmptyOr(f.options.InitDurationIfFail, previousDuration)
+		} else {
+			duration = previousDuration
+		}
+
+		duration = util.EmptyOr(duration, time.Second)
+		if f.options.Factor != 0 {
+			duration = time.Duration(float64(duration) * f.options.Factor)
+		}
+		if f.options.Jitter > 0 {
+			duration = Jitter(duration, f.options.Jitter)
+		}
+		if f.options.MaxDuration > 0 && duration > f.options.MaxDuration {
+			duration = f.options.MaxDuration
+		}
+		return duration
+	}
+	return f.options.Duration
+}
+
+func BackoffUntil(f func() error, backoff BackoffManager, sliding bool, stopCh <-chan struct{}) {
+	var delay time.Duration
+	previousError := false
+
+	ticker := time.NewTicker(backoff.Backoff(delay, previousError))
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-stopCh:
+			return
+		default:
+		}
+
+		if !sliding {
+			delay = backoff.Backoff(delay, previousError)
+		}
+
+		if err := f(); err != nil {
+			previousError = true
+		} else {
+			previousError = false
+		}
+
+		if sliding {
+			delay = backoff.Backoff(delay, previousError)
+		}
+
+		ticker.Reset(delay)
+		select {
+		case <-stopCh:
+			return
+		default:
+		}
+
+		select {
+		case <-stopCh:
+			return
+		case <-ticker.C:
+		}
+	}
+}
+
+// Jitter returns a time.Duration between duration and duration + maxFactor *
+// duration.
+//
+// This allows clients to avoid converging on periodic behavior. If maxFactor
+// is 0.0, a suggested default value will be chosen.
+func Jitter(duration time.Duration, maxFactor float64) time.Duration {
+	if maxFactor <= 0.0 {
+		maxFactor = 1.0
+	}
+	wait := duration + time.Duration(rand.Float64()*maxFactor*float64(duration))
+	return wait
+}
+
+func Until(f func(), period time.Duration, stopCh <-chan struct{}) {
+	ff := func() error {
+		f()
+		return nil
+	}
+	BackoffUntil(ff, BackoffFunc(func(time.Duration, bool) time.Duration {
+		return period
+	}), true, stopCh)
+}
+
+func MergeAndCloseOnAnyStopChannel[T any](upstreams ...<-chan T) <-chan T {
+	out := make(chan T)
+
+	for _, upstream := range upstreams {
+		ch := upstream
+		go lo.Try0(func() {
+			select {
+			case <-ch:
+				close(out)
+			case <-out:
+			}
+		})
+	}
+	return out
+}