一次退出所有递归生成的goroutine

Question

I have a function that recursively spawns goroutines to walk a DOM tree, putting the nodes they find into a channel shared between all of them.

import (
    "golang.org/x/net/html"
    "sync"
)

func walk(doc *html.Node, ch chan *html.Node) {
    var wg sync.WaitGroup
    defer close(ch)
    var f func(*html.Node)
    f = func(n *html.Node) {
        defer wg.Done()
        ch <- n
        for c := n.FirstChild; c != nil; c = c.NextSibling {
            wg.Add(1)
            go f(c)
        }
    }
    wg.Add(1)
    go f(doc)
    wg.Wait()
}

Which I'd call like

// get the webpage using http
// parse the html into doc
ch := make(chan *html.Node)
go walk(doc, ch)

for c := range ch {
    if someCondition(c) {
        // do something with c
        // quit all goroutines spawned by walk
    }
}

I am wondering how I could quit all of these goroutines--i.e. close ch--once I have found a node of a certain type or some other condition has been fulfilled. I have tried using a quit channel that'd be polled before spawning the new goroutines and close ch if a value was received but that lead to race conditions where some goroutines tried sending on the channel that had just been closed by another one. I was pondering using a mutex but it seems inelegant and against the spirit of go to protect a channel with a mutex. Is there an idiomatic way to do this using channels? If not, is there any way at all? Any input appreciated!

Answer 1

The context package provides similar functionality. Using context.Context with a few Go-esque patterns, you can achieve what you need.

To start you can check this article to get a better feel of cancellation with context: https://www.sohamkamani.com/blog/golang/2018-06-17-golang-using-context-cancellation/

Also make sure to check the official GoDoc: https://golang.org/pkg/context/

So to achieve this functionality your function should look more like:

func walk(ctx context.Context, doc *html.Node, ch chan *html.Node) {
    var wg sync.WaitGroup
    defer close(ch)

    var f func(*html.Node)
    f = func(n *html.Node) {
        defer wg.Done()

        ch <- n
        for c := n.FirstChild; c != nil; c = c.NextSibling {
            select {
            case <-ctx.Done():
                return // quit the function as it is cancelled
            default:
                wg.Add(1)
                go f(c)
            }
        }
    }

    select {
    case <-ctx.Done():
        return // perhaps it was cancelled so quickly
    default:
        wg.Add(1)
        go f(doc)
        wg.Wait()
    }
}

And when calling the function, you will have something like:

// ...
ctx, cancelFunc := context.WithCancel(context.Background())
walk(ctx, doc, ch)
for value := range ch {
    // ...
    if someCondition {
        cancelFunc()
        // the for loop will automatically exit as the channel is being closed for the inside
    }
}