feat: implement cancellation

I've also added 1., 2., etc. to the printed list of URLs.

1 files changed, 20 insertions, 3 deletions

diff --git a/main.go b/main.go
index ed203d5..d23b4cf 100644
--- a/main.go
+++ b/main.go
@@ -1,6 +1,7 @@
 package main
 
 import (
+	"context"
 	"flag"
 	"fmt"
 	"log"
@@ -33,26 +34,42 @@ func main() {
 	}()
 
 	sem := make(chan struct{}, *maxConcurrency)
+	ctx, cancel := context.WithCancel(context.Background())
 
 	seen := make(map[url.URL]bool)
+	i := 1
+
+loop:
 	for list := range worklist {
 		for _, u := range list {
+			if len(seen) == 30 {
+				cancel()
+				break loop
+			}
+
 			if !seen[u] {
-				fmt.Printf("%s\n", &u)
+				fmt.Printf("%d. %s\n", i, &u)
+				i++
 				seen[u] = true
 
 				go func() {
 					sem <- struct{}{}
 					defer func() { <-sem }()
 
-					worklist <- crawl(u)
+					worklist <- crawl(ctx, u)
 				}()
 			}
 		}
 	}
 }
 
-func crawl(u url.URL) []url.URL {
+func crawl(ctx context.Context, u url.URL) []url.URL {
+	select {
+	case <-ctx.Done():
+		return nil
+	default:
+	}
+
 	doc, err := fetch(u)
 	if err != nil {
 		log.Print(err)