feat: add early termination condition based on maxURLs

1 files changed, 10 insertions, 1 deletions

diff --git a/workers.go b/workers.go
index fcfb504..ce1001e 100644
--- a/workers.go
+++ b/workers.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"net/url"
 	"sync"
@@ -24,13 +25,20 @@ func workers(startURL url.URL, maxConcurrency, maxURLs int) {
 	}()
 
 	var wg sync.WaitGroup
+	ctx, cancel := context.WithCancel(context.Background())
+
 	// Create maxConcurrency worker goroutines to demultiplex from
 	// the urls channel (unseen links.)
 	for range maxConcurrency {
 		wg.Go(func() {
 			for u := range urls {
 				batch := getBatch(u)
-				go func() { worklist <- batch }()
+				select {
+				case <-ctx.Done():
+					return
+				default:
+					go func() { worklist <- batch }()
+				}
 			}
 		})
 	}
@@ -58,5 +66,6 @@ loop:
 		}
 	}
 
+	cancel()
 	wg.Wait()
 }