Merge pull request #14 from coreweave/rwang.fixs3chan.10222023

Fix channels for s3
coreweave · Oct 25, 2023 · eb30a00 · eb30a00
2 parents 1bca3fc + 7994a99
commit eb30a00
Showing 1 changed file with 10 additions and 8 deletions.
diff --git a/cmd/dataset_tokenizer/dataset_tokenizer.go b/cmd/dataset_tokenizer/dataset_tokenizer.go
@@ -416,17 +416,19 @@ func ReadTextsFromS3(
 		go startReader()
 	}
 
-	// List objects recursively.
-	getObjectsS3Recursively(svc, bucketName, "", objects)
+	go func() {
+		// List objects recursively.
+		getObjectsS3Recursively(svc, bucketName, "", objects)
 
-	// Close the objects channel when done.
-	close(objects)
+		// Close the objects channel when done.
+		close(objects)
 
-	// Wait for all reader goroutines to finish.
-	wg.Wait()
+		// Wait for all reader goroutines to finish.
+		wg.Wait()
 
-	// Close the runeReaders channel.
-	close(runeReaders)
+		// Close the runeReaders channel.
+		close(runeReaders)
+	}()
 
 	return runeReaders, nil
 }