golang 协程扫描的退出问题

在进行协程扫描网络端口时，程序可能会提前退出。这是什么原因造成的呢？

问题解析

其主要问题在于程序没有正确保证在关闭管道前，协程已经完成所有逻辑处理。在给定的代码中，主协程使用了 time.sleep(time.second) 来休眠，试图等待协程处理完工作。然而，这种做法并不能可靠地确保协程已完成工作。若由于扫描较慢或其他原因，导致协程在休眠时间内仍未处理完工作，则主协程将关闭管道，导致协程无法写入结果。

立即学习“go语言免费学习笔记（深入）”；

解决方案

Golang Uber Dig

使用 uber-go/dig 在 Golang 中实现依赖注入——基于反射容器，提供 Provide/Invoke、dig.In/dig.Out 参数与结果对象，支持命名依赖等

下载

为了解决这个问题，需要使用同步机制，例如 sync.waitgroup 来等待协程正确完成工作。具体做法如下：

• 在主协程中创建 sync.waitgroup 实例。
• 在每个协程的开头调用 wg.add(1) 来增加waitgroup计数。
• 在协程的最后调用 wg.done() 来减少waitgroup计数。
• 在主协程中调用 wg.wait() 来等待waitgroup计数变为0，表示所有协程已完成工作。

如此一来，主协程可以可靠地等待协程完成扫描，再关闭管道。

代码改善

改进后的代码如下：

package main

import (
    "fmt"
    "net"
    "sync"
    "time"
)

func scanPort(jobChan <-chan int, retChan chan<- string, wg *sync.WaitGroup) {
    defer wg.Done()
    // 从管道取端口
    for port := range jobChan {
        // 发起tcp连接
        target := fmt.Sprintf("149.129.68.235:%d", port)
        //fmt.Println(target)
        _, err := net.DialTimeout("tcp", target, 100*time.Millisecond)
        //_, err := net.Dial("tcp", target)
        // 未发生错误时写入结果管道
        if err == nil {
            retChan <- fmt.Sprintf("%s is open", target)
        }

    }
}

func main() {
    jobChan := make(chan int, 5)
    retChan := make(chan string, 5)

    // 写入1024个待扫描端口到任务管道
    go func() {
        for i := 0; i < 1024; i++ {
            jobChan <- i
        }
        close(jobChan)
    }()
    wg := sync.WaitGroup{}
    // 创建n个协程
    for i := 0; i < 10; i++ {
        wg.Add(1)
        go scanPort(jobChan, retChan, &wg)
    }
    go func() {
        for ret := range retChan {
            fmt.Println(ret)
        }
    }()
    // 关闭
    wg.Wait()
    close(retChan)
}

现在，主协程将等待 waitgroup 计数变为0，表示所有协程已完成工作，然后再关闭管道，确保数据被正确写入。