fix: make sure all workers terminate by the end

1 files changed, 14 insertions, 3 deletions

diff --git a/workers.go b/workers.go
index ce1001e..01c6888 100644
--- a/workers.go
+++ b/workers.go
@@ -29,17 +29,26 @@ func workers(startURL url.URL, maxConcurrency, maxURLs int) {
 
 	// Create maxConcurrency worker goroutines to demultiplex from
 	// the urls channel (unseen links.)
-	for range maxConcurrency {
+	for i := range maxConcurrency {
 		wg.Go(func() {
+		loop:
 			for u := range urls {
 				batch := getBatch(u)
 				select {
 				case <-ctx.Done():
-					return
+					break loop
 				default:
-					go func() { worklist <- batch }()
+					go func() {
+						select {
+						case <-ctx.Done():
+							return
+						case worklist <- batch:
+						}
+					}()
 				}
 			}
+
+			fmt.Printf("terminating %d\n", i+1)
 		})
 	}
 
@@ -66,6 +75,8 @@ loop:
 		}
 	}
 
+	close(urls)
+
 	cancel()
 	wg.Wait()
 }