WaitGroup

Hello!

Today we will learn what is WaitGroup in golang and we will look how to apply them on an example of the program for files search.

WaitGroup is waiting for the completion of go routines. The main go routine calls the Add method, and then each go routine calls the Done method when the routine ends.

Let’s start with the findFile method. In this method we read all files and directories in the root folder. We loop through them and if the name of the file what we are looking for matches the name of the file from the loop then we add it to the list of matches. If it is a directory, then recursively call a search in it

package main

import (
    "fmt"
    "io/ioutil"
    "path/filepath"
    "strings"
)

var (
    matches []string
)

func findFile(root string, filename string) {
    fmt.Println("Searching in", root)
    files, _ := ioutil.ReadDir(root)
    for _, file := range files {
        if strings.Contains(file.Name(), filename) {
            matches = append(matches, filepath.Join(root, file.Name()))
        }
        if file.IsDir() {
            findFile(filepath.Join(root, file.Name()), filename)
        }
    }
}

func main() {
    findFile("/var/", "secret.txt")
    for _, file := range matches {
        fmt.Println("Matched", file)
    }
}

The search works, but it takes a long time. You can speed it up with go routines. To do this, add the word go before calling the findFile method and before the recursive call. But in this case, the search will not work because the script will exit immediately, because there is nothing to wait for the go routines. To fix this WaitGroup can be used. The script will only complete when all go routines call the Done method.

package main

import (
    "fmt"
    "io/ioutil"
    "path/filepath"
    "strings"
    "sync"
)

var (
    matches   []string
    waitGroup = sync.WaitGroup{}
    lock      = sync.Mutex{}
)

func findFile(root string, filename string) {
    fmt.Println("Searching in", root)
    files, _ := ioutil.ReadDir(root)
    for _, file := range files {
        if strings.Contains(file.Name(), filename) {
            lock.Lock()
            matches = append(matches, filepath.Join(root, file.Name()))
            lock.Unlock()
        }
        if file.IsDir() {
            waitGroup.Add(1)
            go findFile(filepath.Join(root, file.Name()), filename)
        }
    }
    waitGroup.Done()
}

func main() {
    waitGroup.Add(1)
    go findFile("/var/", "secret.txt")
    waitGroup.Wait()
    for _, file := range matches {
        fmt.Println("Matched", file)
    }
}

A few things have been added here. In the findFile method, lock was added before adding file names to list. To ensure that routines do not overwrite data. Another thing that has been added is waitGroup. Before the each go routine call, the waitGroup.Add(1) call is added, and after the go routine, waitGroup.Done() is called. And to the main method added waitGroup.Wait() which will wait while all go routines finish work. And then a list of all files will be displayed.