循环中的 GO 例程 - 函数调用

Question

问题描述

我想向我的for 循环 中的每个迭代发送一个并发调用。

当前代码

这个程序主要是尝试在服务器上找到一个有效的目录

package main

import (
    "bufio"
    "fmt"
    "net/http"
    "os"
)

func main() {
    var site string
    if len(os.Args) > 1 {
        site = os.Args[1]
    } else {
        fmt.Println("Need the host. Example > http://www.google.com/")
        fmt.Println("wfuzz http://www.google.com/ wordlist.txt")
        site = "http://www.google.com"
        os.Exit(1)
    }
    f, err := os.Open(os.Args[2])

    if err != nil {
        fmt.Println("error opening file!", err)
    }
    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        a := site + scanner.Text()
        get_page(a)
    }
    if err := scanner.Err(); err != nil {
        fmt.Println("error", err)
    }

    var input string
    fmt.Scanln(&input)
    fmt.Println("done")

}

func get_page(site string) string {
    a := "\n"
    res, err := http.Get(site)
    if err != nil {
        fmt.Println("error:", err)
    }
    if res.StatusCode != 404 {
        fmt.Println("Page Found!! +::", site, " ::+", res.StatusCode)
    }
    return a
}

要实现的目标

我想要实现的主要是发送我的 for 循环并发调用，这样我会加快进程，这将是当前程序的最终目标。

Answer 1

结合使用 Go 标准库 sync.WaitGroup 和 channel ，您可以完成此操作，而无需明确跟踪要读取的文件中的单词数。

修改您的 get_page(旁注:转到函数名称中的用户下划线不是惯用的 -> https://golang.org/doc/effective_go.html#names )函数以接受额外的 channel 和 WaitGroup ，您可以执行类似以下:

package main

import (
    "bufio"
    "fmt"
    "net/http"
    "os"
    "sync"
)

func main() {
    
    //your previous code for getting the command line arguments from os.Args, then...

    //create a WaitGroup
    wg := new(sync.WaitGroup)

    //create a channel that you can write your results too
    resultChan := make(chan string)

    scanner := bufio.NewScanner(f)
    for scanner.Scan() {
        //increment the waitgroup
        wg.Add(1)
        a := site + scanner.Text()
        //call the getPage with the http target, with the result chan, and the waitgroup
        go getPage(a, resultChan, wg)
    }

    if err := scanner.Err(); err != nil {
        fmt.Println("error", err)
        //if there is an error, close the result chan so progam can exit (for more robust cancellation of potential currently running go routines, take a look at https://blog.golang.org/pipelines)
        close(resultChan)
    }

    //start another go routine to wait for all concurrent calls to get_page to complete, then close the result channel
    go func() {
        wg.Wait()
        close(resultChan)
    }()

    for result := range resultChan {
        //deal with the result in some way
        fmt.Println(result)
    }

    fmt.Println("done")
}

//side note - it is not idiomatic go to use underscores in func name
func getPage(site string, ch chan<- string, wg *sync.WaitGroup) {
    //decrement waitgroup when operation completes
    defer wg.Done()

    res, err := http.Get(site)
    if err != nil {
        ch <- fmt.Sprintf("error: %s", err)
        return
    }
    if res.StatusCode != 404 {
        ch <- fmt.Sprintf("Page Found!! +:: %s ::+ %d", site, res.StatusCode)
    }
}