feat: added more control over the transport layer

Velka-DEV · May 28, 2024 · 2044a09 · 2044a09
1 parent 920c5ed
commit 2044a09
Show file tree

Hide file tree

Showing 3 changed files with 49 additions and 20 deletions.
diff --git a/pkg/core/builder.go b/pkg/core/builder.go
@@ -1,23 +1,45 @@
 package core
 
+import "time"
+
+type TransportConfig struct {
+	DisableKeepAlive    bool
+	MaxIdleConns        int
+	MaxIdleConnsPerHost int
+	IdleConnTimeout     time.Duration
+}
+
 type CheckerBuilder struct {
-	config *CheckerConfig
-	queue  *ComboQueue
-	pool   *HTTPClientPool
+	config          *CheckerConfig
+	transportConfig *TransportConfig
+	queue           *ComboQueue
+	pool            *HTTPClientPool
 
 	checkProcess  CheckProcess
 	outputProcess OutputProcess
 }
 
 func NewCheckerBuilder() *CheckerBuilder {
-	return &CheckerBuilder{}
+	return &CheckerBuilder{
+		transportConfig: &TransportConfig{
+			DisableKeepAlive:    false,
+			MaxIdleConns:        100,
+			MaxIdleConnsPerHost: 2,
+			IdleConnTimeout:     90 * time.Second,
+		},
+	}
 }
 
 func (b *CheckerBuilder) WithConfig(config *CheckerConfig) *CheckerBuilder {
 	b.config = config
 	return b
 }
 
+func (b *CheckerBuilder) WithTransportConfig(transportConfig *TransportConfig) *CheckerBuilder {
+	b.transportConfig = transportConfig
+	return b
+}
+
 func (b *CheckerBuilder) WithCombos(combos []*Combo) *CheckerBuilder {
 
 	if len(combos) == 0 {
@@ -29,7 +51,7 @@ func (b *CheckerBuilder) WithCombos(combos []*Combo) *CheckerBuilder {
 }
 
 func (b *CheckerBuilder) WithProxies(proxies []string, scheme string) *CheckerBuilder {
-	b.pool = NewHTTPClientPoolWithProxies(proxies, scheme, b.config.DisableKeepAlive)
+	b.pool = NewHTTPClientPoolWithProxies(proxies, scheme, b.transportConfig)
 	return b
 }
 
@@ -57,7 +79,7 @@ func (b *CheckerBuilder) Build() *Checker {
 	}
 
 	if b.pool == nil {
-		b.pool = NewHTTPClientPool(b.config.DisableKeepAlive)
+		b.pool = NewHTTPClientPool(b.transportConfig)
 	}
 
 	return NewChecker(

diff --git a/pkg/core/checker.go b/pkg/core/checker.go
@@ -11,15 +11,14 @@ import (
 )
 
 type CheckerConfig struct {
-	Threads          int
-	DisableKeepAlive bool
-	WaitBeforeRetry  time.Duration
-	OutputToFile     bool
-	OutputDirectory  string
-	OutputFree       bool
-	OutputInvalid    bool
-	OutputLocked     bool
-	OutputUnknown    bool
+	Threads         int
+	WaitBeforeRetry time.Duration
+	OutputToFile    bool
+	OutputDirectory string
+	OutputFree      bool
+	OutputInvalid   bool
+	OutputLocked    bool
+	OutputUnknown   bool
 }
 
 type CheckProcessArgs struct {

diff --git a/pkg/core/httppool.go b/pkg/core/httppool.go
@@ -17,13 +17,18 @@ type HTTPClientPool struct {
 	index   int
 }
 
-func NewHTTPClientPool(disableKeepAlive bool) *HTTPClientPool {
+func NewHTTPClientPool(config *TransportConfig) *HTTPClientPool {
 	return &HTTPClientPool{clients: []ClientWithProxy{{Client: &http.Client{
-		Transport: &http.Transport{DisableKeepAlives: disableKeepAlive},
+		Transport: &http.Transport{
+			DisableKeepAlives:   config.DisableKeepAlive,
+			MaxIdleConns:        config.MaxIdleConns,
+			MaxIdleConnsPerHost: config.MaxIdleConnsPerHost,
+			IdleConnTimeout:     config.IdleConnTimeout,
+		},
 	}, Proxy: ""}}}
 }
 
-func NewHTTPClientPoolWithProxies(proxies []string, scheme string, disableKeepAlive bool) *HTTPClientPool {
+func NewHTTPClientPoolWithProxies(proxies []string, scheme string, config *TransportConfig) *HTTPClientPool {
 	if len(proxies) == 0 {
 		panic("proxies are empty")
 	}
@@ -38,8 +43,11 @@ func NewHTTPClientPoolWithProxies(proxies []string, scheme string, disableKeepAl
 
 		client := &http.Client{
 			Transport: &http.Transport{
-				Proxy:             http.ProxyURL(proxyURL),
-				DisableKeepAlives: disableKeepAlive,
+				Proxy:               http.ProxyURL(proxyURL),
+				DisableKeepAlives:   config.DisableKeepAlive,
+				MaxIdleConns:        config.MaxIdleConns,
+				MaxIdleConnsPerHost: config.MaxIdleConnsPerHost,
+				IdleConnTimeout:     config.IdleConnTimeout,
 			},
 		}
 		clients = append(clients, ClientWithProxy{Client: client, Proxy: proxy})