custom dialer only for high thread mode

Author: tanq

cmd/root.go

@@ -38,12 +38,15 @@ var rootCmd = &cobra.Command{
 		if len(args) == 0 && urlListFile == "" {
 			log.Fatal().Msg("No URL or URL list provided")
 		}
-		url := args[0]
-		if _, err := u.Parse(url); err != nil {
-			log.Fatal().Err(err).Msg("Invalid URL format")
+		if urlListFile != "" && len(args) > 0 {
+			log.Fatal().Msg("Cannot specify url argument and --urllist together, choose one")
 		}
-		if urlListFile != "" && url != "" {
-			log.Fatal().Msg("Cannot specify both argument and --urllist, choose one")
+		url := ""
+		if len(args) > 1 {
+			url = args[0]
+			if _, err := u.Parse(url); err != nil {
+				log.Fatal().Err(err).Msg("Invalid URL format")
+			}
 		}
 
 		// Handle single URL download

internal/downloader.go

@@ -2,6 +2,7 @@ package internal
 
 import (
 	"fmt"
+	"net/http"
 	"os"
 	"path/filepath"
 	"sync"
@@ -51,11 +52,11 @@ func BatchDownload(entries []DownloadEntry, numLinks int, connectionsPerLink int
 					ProxyURL:    proxyURL,
 				}
 				progressCh := make(chan int64)
-				client := createHTTPClient(config.Timeout, config.KATimeout, config.ProxyURL)
+				useHighThreadMode := config.Connections > 6
+				client := createHTTPClient(config.Timeout, config.KATimeout, config.ProxyURL, useHighThreadMode)
 				fileSize, err := getFileSize(config.URL, config.UserAgent, client)
 
 				if err == ErrRangeRequestsNotSupported {
-					// file size unknown, so can't show progress; so track bytes downloaded
 					logger.Debug().Str("url", entry.URL).Msg("Range requests not supported, using simple download")
 					progressManager.Register(entry.OutputPath, -1) // -1 means unknown size
 				} else if err != nil {
@@ -82,9 +83,8 @@ func BatchDownload(entries []DownloadEntry, numLinks int, connectionsPerLink int
 				if err == ErrRangeRequestsNotSupported {
 					err = performSimpleDownload(entry.URL, entry.OutputPath, client, config.UserAgent, progressCh)
 				} else {
-					err = downloadWithProgress(config, fileSize, progressCh)
+					err = downloadWithProgress(config, client, fileSize, progressCh)
 				}
-				// close(progressCh)
 				progressWg.Wait()
 
 				if err != nil {
@@ -110,9 +110,9 @@ func BatchDownload(entries []DownloadEntry, numLinks int, connectionsPerLink int
 	return nil
 }
 
-func downloadWithProgress(config DownloadConfig, fileSize int64, progressCh chan<- int64) error {
+func downloadWithProgress(config DownloadConfig, client *http.Client, fileSize int64, progressCh chan<- int64) error {
 	log := GetLogger("download-worker")
-	client := createHTTPClient(config.Timeout, config.KATimeout, config.ProxyURL)
+	// client := createHTTPClient(config.Timeout, config.KATimeout, config.ProxyURL)
 	log.Debug().Str("size", formatBytes(uint64(fileSize))).Msg("File size determined")
 	job := DownloadJob{
 		Config:    config,

internal/helpers.go

@@ -15,7 +15,7 @@ import (
 	"gopkg.in/yaml.v3"
 )
 
-const bufferSize = 1024 * 1024 * 2 // 2MB buffer
+const bufferSize = 1024 * 1024 * 8 // 8MB buffer
 const ToolUserAgent = "danzo/1337"
 
 var chunkIDRegex = regexp.MustCompile(`\.part(\d+)$`)
@@ -108,7 +108,7 @@ func ReadDownloadList(filePath string) ([]DownloadEntry, error) {
 	return entries, nil
 }
 
-func createHTTPClient(timeout time.Duration, keepAliveTO time.Duration, proxyURL string) *http.Client {
+func createHTTPClient(timeout time.Duration, keepAliveTO time.Duration, proxyURL string, highThreadMode bool) *http.Client {
 	transport := &http.Transport{
 		MaxIdleConns:        100,
 		MaxIdleConnsPerHost: 100, // for connection reuse
@@ -118,7 +118,9 @@ func createHTTPClient(timeout time.Duration, keepAliveTO time.Duration, proxyURL
 		// These two seem to reduce performance drastically with custom dial context
 		// DisableKeepAlives:   false,
 		// ForceAttemptHTTP2:   true,
-		DialContext: (&net.Dialer{
+	}
+	if highThreadMode {
+		transport.DialContext = (&net.Dialer{
 			Timeout:   30 * time.Second,
 			KeepAlive: 30 * time.Second,
 			DualStack: true,
@@ -128,7 +130,7 @@ func createHTTPClient(timeout time.Duration, keepAliveTO time.Duration, proxyURL
 					setSocketOptions(fd)
 				})
 			},
-		}).DialContext,
+		}).DialContext
 	}
 	if proxyURL != "" {
 		proxyURLParsed, err := url.Parse(proxyURL)

internal/operations.go

@@ -59,37 +59,6 @@ func parseContentLength(contentLength string) (int64, error) {
 	return size, nil
 }
 
-// func getFileSize(url string, userAgent string, client *http.Client) (int64, error) {
-// 	log := GetLogger("filesize")
-// 	req, err := http.NewRequest("HEAD", url, nil)
-// 	if err != nil {
-// 		return 0, err
-// 	}
-// 	req.Header.Set("User-Agent", userAgent)
-// 	log.Debug().Str("url", url).Msg("Sending HEAD request")
-// 	resp, err := client.Do(req)
-// 	if err != nil {
-// 		return 0, err
-// 	}
-// 	defer resp.Body.Close()
-// 	if resp.Header.Get("Accept-Ranges") != "bytes" {
-// 		return 0, errors.New("server doesn't support range requests")
-// 	}
-// 	contentLength := resp.Header.Get("Content-Length")
-// 	if contentLength == "" {
-// 		return 0, errors.New("server didn't provide Content-Length header")
-// 	}
-// 	size, err := strconv.ParseInt(contentLength, 10, 64)
-// 	if err != nil {
-// 		return 0, fmt.Errorf("invalid content length: %v", err)
-// 	}
-// 	if size <= 0 {
-// 		return 0, errors.New("invalid file size reported by server")
-// 	}
-// 	log.Debug().Int64("bytes", size).Msg("File size determined")
-// 	return size, nil
-// }
-
 func downloadChunk(job *DownloadJob, chunk *DownloadChunk, client *http.Client, wg *sync.WaitGroup, progressCh chan<- int64, mutex *sync.Mutex) {
 	log := GetLogger("chunk").With().Int("chunkId", chunk.ID).Logger()
 	defer wg.Done()

internal/simple-downloader.go

@@ -19,7 +19,7 @@ func SimpleDownload(url string, outputPath string) error {
 		progressManager.ShowSummary()
 	}()
 
-	client := createHTTPClient(3*time.Minute, 90*time.Second, "")
+	client := createHTTPClient(3*time.Minute, 90*time.Second, "", false)
 	progressCh := make(chan int64)
 
 	req, err := http.NewRequest("HEAD", url, nil)

internal/socket-unix.go

@@ -7,6 +7,9 @@ import (
 )
 
 func setSocketOptions(fd uintptr) {
-	syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_RCVBUF, 1024*1024)
-	syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_SNDBUF, 1024*1024)
+	syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_NODELAY, 1) // Disable Nagle's algorithm
+	// syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, syscall.TCP_WINDOW_CLAMP, 2*bufferSize)
+	// syscall.SetsockoptString(int(fd), syscall.IPPROTO_TCP, syscall.TCP_CONGESTION, "cubic")
+	syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_RCVBUF, bufferSize)
+	syscall.SetsockoptInt(int(fd), syscall.SOL_SOCKET, syscall.SO_SNDBUF, bufferSize)
 }

internal/socket-windows.go

@@ -7,6 +7,7 @@ import (
 )
 
 func setSocketOptions(fd uintptr) {
-	syscall.SetsockoptInt(syscall.Handle(fd), syscall.SOL_SOCKET, syscall.SO_RCVBUF, 1024*1024)
-	syscall.SetsockoptInt(syscall.Handle(fd), syscall.SOL_SOCKET, syscall.SO_SNDBUF, 1024*1024)
+	syscall.SetsockoptInt(syscall.Handle(fd), syscall.IPPROTO_TCP, syscall.TCP_NODELAY, 1) // Disable Nagle's algorithm
+	syscall.SetsockoptInt(syscall.Handle(fd), syscall.SOL_SOCKET, syscall.SO_RCVBUF, bufferSize)
+	syscall.SetsockoptInt(syscall.Handle(fd), syscall.SOL_SOCKET, syscall.SO_SNDBUF, bufferSize)
 }