feat: add cancellation feature

1 files changed, 19 insertions, 3 deletions

diff --git a/classic.go b/classic.go
index 9adb5bf..ef138e2 100644
--- a/classic.go
+++ b/classic.go
@@ -1,9 +1,11 @@
 package main
 
 import (
+	"context"
 	"fmt"
 	"log"
 	"net/url"
+	"sync"
 )
 
 func classic(startURL url.URL, maxConcurrency, maxURLs int) {
@@ -19,6 +21,10 @@ func classic(startURL url.URL, maxConcurrency, maxURLs int) {
 	seen := make(map[url.URL]bool)
 	count := 1
 
+	ctx, cancel := context.WithCancel(context.Background())
+	var wg sync.WaitGroup
+
+loop:
 	for ; numPendingSends > 0; numPendingSends-- {
 		batch := <-worklist
 		for _, u := range batch {
@@ -27,14 +33,24 @@ func classic(startURL url.URL, maxConcurrency, maxURLs int) {
 				count++
 
 				seen[u] = true
+				if len(seen) == maxURLs {
+					break loop
+				}
 
 				numPendingSends++
-				go func() {
-					worklist <- getBatch(u)
-				}()
+				wg.Go(func() {
+					select {
+					case <-ctx.Done():
+						return
+					case worklist <- getBatch(u):
+					}
+				})
 			}
 		}
 	}
+
+	cancel()
+	wg.Wait()
 }
 
 func getBatch(u url.URL) []url.URL {