feat: support customize transport in daemon (#866)

Signed-off-by: Jim Ma <[email protected]>
dragonflyoss · Dec 3, 2021 · 4caae06 · 4caae06
1 parent 3e4a047
commit 4caae06
Show file tree

Hide file tree

Showing 5 changed files with 62 additions and 3 deletions.
diff --git a/client/config/peerhost.go b/client/config/peerhost.go
@@ -170,6 +170,17 @@ type DownloadOption struct {
 	DownloadGRPC         ListenOption         `mapstructure:"downloadGRPC" yaml:"downloadGRPC"`
 	PeerGRPC             ListenOption         `mapstructure:"peerGRPC" yaml:"peerGRPC"`
 	CalculateDigest      bool                 `mapstructure:"calculateDigest" yaml:"calculateDigest"`
+	TransportOption      *TransportOption     `mapstructure:"transportOption" yaml:"transportOption"`
+}
+
+type TransportOption struct {
+	DialTimeout           time.Duration `mapstructure:"dialTimeout" yaml:"dialTimeout"`
+	KeepAlive             time.Duration `mapstructure:"keepAlive" yaml:"keepAlive"`
+	MaxIdleConns          int           `mapstructure:"maxIdleConns" yaml:"maxIdleConns"`
+	IdleConnTimeout       time.Duration `mapstructure:"idleConnTimeout" yaml:"idleConnTimeout"`
+	ResponseHeaderTimeout time.Duration `mapstructure:"responseHeaderTimeout" yaml:"responseHeaderTimeout"`
+	TLSHandshakeTimeout   time.Duration `mapstructure:"tlsHandshakeTimeout" yaml:"tlsHandshakeTimeout"`
+	ExpectContinueTimeout time.Duration `mapstructure:"expectContinueTimeout" yaml:"expectContinueTimeout"`
 }
 
 type ProxyOption struct {

diff --git a/client/daemon/daemon.go b/client/daemon/daemon.go
@@ -158,7 +158,7 @@ func New(opt *config.DaemonOption) (Daemon, error) {
 	pieceManager, err := peer.NewPieceManager(storageManager,
 		opt.Download.PieceDownloadTimeout,
 		peer.WithLimiter(rate.NewLimiter(opt.Download.TotalRateLimit.Limit, int(opt.Download.TotalRateLimit.Limit))),
-		peer.WithCalculateDigest(opt.Download.CalculateDigest),
+		peer.WithCalculateDigest(opt.Download.CalculateDigest), peer.WithTransportOption(opt.Download.TransportOption),
 	)
 	if err != nil {
 		return nil, err

diff --git a/client/daemon/peer/piece_manager.go b/client/daemon/peer/piece_manager.go
@@ -20,6 +20,8 @@ import (
 	"context"
 	"io"
 	"math"
+	"net"
+	"net/http"
 	"time"
 
 	"golang.org/x/time/rate"
@@ -76,13 +78,43 @@ func WithCalculateDigest(enable bool) func(*pieceManager) {
 	}
 }
 
-// WithLimiter sets upload rate limiter, the burst size must big than piece size
+// WithLimiter sets upload rate limiter, the burst size must be bigger than piece size
 func WithLimiter(limiter *rate.Limiter) func(*pieceManager) {
 	return func(manager *pieceManager) {
 		manager.Limiter = limiter
 	}
 }
 
+func WithTransportOption(opt *config.TransportOption) func(*pieceManager) {
+	return func(manager *pieceManager) {
+		if opt == nil {
+			return
+		}
+		if opt.IdleConnTimeout > 0 {
+			defaultTransport.(*http.Transport).IdleConnTimeout = opt.IdleConnTimeout
+		}
+		if opt.DialTimeout > 0 && opt.KeepAlive > 0 {
+			defaultTransport.(*http.Transport).DialContext = (&net.Dialer{
+				Timeout:   opt.DialTimeout,
+				KeepAlive: opt.KeepAlive,
+				DualStack: true,
+			}).DialContext
+		}
+		if opt.MaxIdleConns > 0 {
+			defaultTransport.(*http.Transport).MaxIdleConns = opt.MaxIdleConns
+		}
+		if opt.ExpectContinueTimeout > 0 {
+			defaultTransport.(*http.Transport).ExpectContinueTimeout = opt.ExpectContinueTimeout
+		}
+		if opt.ResponseHeaderTimeout > 0 {
+			defaultTransport.(*http.Transport).ResponseHeaderTimeout = opt.ResponseHeaderTimeout
+		}
+		if opt.TLSHandshakeTimeout > 0 {
+			defaultTransport.(*http.Transport).TLSHandshakeTimeout = opt.TLSHandshakeTimeout
+		}
+	}
+}
+
 func (pm *pieceManager) DownloadPiece(ctx context.Context, pt Task, request *DownloadPieceRequest) (success bool) {
 	var (
 		start = time.Now().UnixNano()

diff --git a/client/daemon/upload/upload_manager.go b/client/daemon/upload/upload_manager.go
@@ -64,7 +64,7 @@ func NewUploadManager(s storage.Manager, opts ...func(*uploadManager)) (Manager,
 	return u, nil
 }
 
-// WithLimiter sets upload rate limiter, the burst size must big than piece size
+// WithLimiter sets upload rate limiter, the burst size must be bigger than piece size
 func WithLimiter(limiter *rate.Limiter) func(*uploadManager) {
 	return func(manager *uploadManager) {
 		manager.Limiter = limiter

diff --git a/docs/en/deployment/configuration/dfget.yaml b/docs/en/deployment/configuration/dfget.yaml
@@ -79,6 +79,22 @@ download:
   perPeerRateLimit: 100Mi
   # download piece timeout
   pieceDownloadTimeout: 30s
+  # golang transport option
+  transportOption:
+    # dial timeout
+    dialTimeout: 2s
+    # keep alive
+    keepAlive: true
+    # same with http.Transport.MaxIdleConns
+    maxIdleConns: 100
+    # same with http.Transport.IdleConnTimeout
+    idleConnTimeout: 90s
+    # same with http.Transport.ResponseHeaderTimeout
+    responseHeaderTimeout: 2s
+    # same with http.Transport.TLSHandshakeTimeout
+    tlsHandshakeTimeout: 1s
+    # same with http.Transport.ExpectContinueTimeout
+    expectContinueTimeout: 2s
   # download grpc option
   downloadGRPC:
     # security option