dkforest

commit 80284583338b2dd496688fb4307a79985a86f52b
parent 529d85fc3527f7b64355793d39e6a5a62a201f01
Author: n0tr1v <n0tr1v@protonmail.com>
Date:   Mon, 30 Jan 2023 23:18:32 -0800

add timeout

Diffstat:
M
cmd/dkfupload/main.go
 | 
16
++++++++++------

1 file changed, 10 insertions(+), 6 deletions(-)
diff --git a/cmd/dkfupload/main.go b/cmd/dkfupload/main.go
@@ -39,12 +39,14 @@ func main() {
 	var dry bool
 	var isLocal bool
 	var maxChunkSize int64
+	var httpTimeout time.Duration
 	filedropUUIDUsage := "dkf filedrop uuid"
 	fileNameUsage := "file to upload"
 	nbThreadsUsage := "nb threads"
 	nbThreadsDefaultValue := 20
 	chunkSizeUsage := "chunk size"
-	chunkSizeDefaultValue := int64(10 << 20)
+	chunkSizeDefaultValue := int64(5 << 20) // 5MB
+	flag.DurationVar(&httpTimeout, "http-timeout", 5*time.Minute, "http timeout")
 	flag.StringVar(&filedropUUID, "filedrop-uuid", "", filedropUUIDUsage)
 	flag.StringVar(&filedropUUID, "u", "", filedropUUIDUsage)
 	flag.StringVar(&fileName, "file", "", fileNameUsage)
@@ -95,7 +97,8 @@ func main() {
 
 	// Init the filedrop and send metadata about the file
 	if !dry {
-		client := doGetClient(isLocal)
+		log.Println("sending metadata")
+		client := doGetClient(isLocal, httpTimeout)
 		body := url.Values{}
 		body.Set("fileName", fs.Name())
 		body.Set("fileSize", strconv.FormatInt(fileSize, 10))
@@ -119,7 +122,7 @@ func main() {
 	// Start worker threads
 	wg.Add(nbThreads)
 	for i := 0; i < nbThreads; i++ {
-		go work(i, wg, chunksCh, isLocal, dry, maxChunkSize, f, baseUrl, filedropUUID)
+		go work(i, wg, chunksCh, isLocal, dry, maxChunkSize, nbChunks, f, baseUrl, filedropUUID, httpTimeout)
 	}
 
 	// Provide worker threads with tasks to do
@@ -136,8 +139,8 @@ func main() {
 	log.Printf("All done\n")
 }
 
-func work(i int, wg *sync.WaitGroup, chunksCh chan int64, isLocal, dry bool, maxChunkSize int64, f *os.File, baseUrl, filedropUUID string) {
-	client := doGetClient(isLocal)
+func work(i int, wg *sync.WaitGroup, chunksCh chan int64, isLocal, dry bool, maxChunkSize, nbChunks int64, f *os.File, baseUrl, filedropUUID string, httpTimeout time.Duration) {
+	client := doGetClient(isLocal, httpTimeout)
 
 	buf := make([]byte, maxChunkSize)
 	for chunkNum := range chunksCh {
@@ -178,7 +181,7 @@ func work(i int, wg *sync.WaitGroup, chunksCh chan int64, isLocal, dry bool, max
 	wg.Done()
 }
 
-func doGetClient(isLocal bool) (client *http.Client) {
+func doGetClient(isLocal bool, httpTimeout time.Duration) (client *http.Client) {
 	hasToSucceed(func() error {
 		if isLocal {
 			client = http.DefaultClient
@@ -192,6 +195,7 @@ func doGetClient(isLocal bool) (client *http.Client) {
 		}
 		return nil
 	})
+	client.Timeout = httpTimeout
 	return
 }