diff --git a/structs.go b/structs.go index 9d4c505..5c0f50e 100644 --- a/structs.go +++ b/structs.go @@ -61,11 +61,13 @@ type Manifest struct { Size int `json:"size"` Offsets []int `json:"offsets"` } `json:"remainder"` - Fragments []struct { - Sha512 string `json:"sha512"` - Size int `json:"size"` - Offsets []int `json:"offsets"` - } `json:"fragments"` + Fragments []Fragment `json:"fragments"` +} + +type Fragment struct { + Sha512 string `json:"sha512"` + Size int `json:"size"` + Offsets []int `json:"offsets"` } type RDBHeader struct { diff --git a/util.go b/util.go index a429438..90cd8d6 100644 --- a/util.go +++ b/util.go @@ -9,6 +9,7 @@ import ( "path/filepath" "slices" "strings" + "sync" ) var Debug = false @@ -112,30 +113,74 @@ func Download(manifest Manifest, downloadDir string, filen string) ([]byte, erro // start downloading chunks ret := make([]byte, manifest.Size) totalMB := float32(manifest.Size) / 1024.0 / 1024.0 - for _, fragment := range manifest.Fragments { - url := transformShaToURL(fragment.Sha512, fragment.Size) - content, err := getContent(url) - if err != nil { - return nil, err - } - downloaded += len(content) * len(fragment.Offsets) - for _, offset := range fragment.Offsets { - trackFrag = append(trackFrag, Frag{Offset: offset, Size: len(content)}) - copy(ret[offset:offset+len(content)], content) - } - if Debug { - // save fragment to file in the fragments/ directory - if err := os.WriteFile(filepath.Join(downloadDir, "fragments", filen, fmt.Sprintf("%d.%s.%d.bytes", fragment.Offsets[0], fragment.Sha512, fragment.Size)), content, os.ModePerm); err != nil { - return nil, err + var wg sync.WaitGroup + var mu sync.Mutex + + fragmentsChan := make(chan Fragment, len(manifest.Fragments)) + errChan := make(chan error, DownThreads) + + // Start worker goroutines + for i := 0; i < DownThreads; i++ { + wg.Add(1) + go func() { + defer wg.Done() + for fragment := range fragmentsChan { + url := transformShaToURL(fragment.Sha512, fragment.Size) + content, err := getContent(url) + if err != nil { + // Handle error, maybe by sending it to an error channel + // For now, we'll just skip and print. + fmt.Printf("Error downloading fragment %s: %v\n", fragment.Sha512, err) + errChan <- err + } + + mu.Lock() + // Safely write to the shared slice and update progress + for _, offset := range fragment.Offsets { + if offset+len(content) > len(ret) { + // Handle case where fragment is too large for the buffer + fmt.Printf("Fragment too large for buffer at offset %d\n", offset) + continue + } + copy(ret[offset:offset+len(content)], content) + } + downloaded += len(content) * len(fragment.Offsets) + + // Print progress + downedMB := float32(downloaded) / 1024.0 / 1024.0 + percent := float32(downloaded) / float32(manifest.Size) * 100.0 + fmt.Printf("\r\033[32mDownloaded\033[0m %.2fMB/%.2fMB (\033[36m%.2f%%\033[0m) of \033[33m%s\033[0m", downedMB, totalMB, percent, filen) + + mu.Unlock() + + // The Debug section can be safely moved here as well + if os.Getenv("DEBUG") == "true" { // Using env var for a cleaner check + mu.Lock() // Re-lock for file write + if err := os.WriteFile(filepath.Join(downloadDir, "fragments", filen, fmt.Sprintf("%d.%s.%d.bytes", fragment.Offsets[0], fragment.Sha512, fragment.Size)), content, os.ModePerm); err != nil { + fmt.Printf("Error saving file: %v\n", err) + } + mu.Unlock() + } } - } - - percent := float32(downloaded) / float32(manifest.Size) * 100.0 - downedMB := float32(downloaded) / 1024.0 / 1024.0 - fmt.Printf("\r\033[32mDownloaded\033[0m %fMB/%fMB (\033[36m%.2f%%\033[0m) of \033[33m%s\033[0m", downedMB, totalMB, percent, filen) + }() } + // Send fragments to the channel + for _, fragment := range manifest.Fragments { + fragmentsChan <- fragment + } + close(fragmentsChan) // Close the channel to signal workers no more tasks are coming + + // Wait for all workers to finish + wg.Wait() + + // Check for errors + if len(errChan) > 0 { + return nil, <-errChan // Return the first error encountered + } + close(errChan) + // sort trackFrag by Offset slices.SortFunc(trackFrag, func(a, b Frag) int { return a.Offset - b.Offset