Golang 协程 TCP 扫描提前退出：如何确保所有协程完成任务？

golang 协程tcp扫描提前退出详解

在进行tcp扫描时，我们期望协程在处理完任务后才退出，然而有时候却出现协程提前退出导致任务未完成的情况。本文对这一问题进行深入分析，给出解决方案。

问题描述

目标是通过协程扫描目标主机范围内的开放端口，例如扫描149.129.68.235的1-1024端口。但代码运行后却发现协程在将结果写入通道之前就离开了。

解决方案

问题在于主协程在协程处理完任务之前就关闭了管道。为了解决这个问题，可以通过使用sync.waitgroup来保证在所有协程完成任务之前，主协程不会关闭管道。

修改后的代码

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 {
        target := fmt.Sprintf("149.129.68.235:%d", port)
        _, err := net.DialTimeout("tcp", target, 100*time.Millisecond)
        if err == nil {
            retChan <- fmt.Sprintf("%s is open", target)
        }
    }
}

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

    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)
}

在修改后的代码中，我们使用sync.waitgroup来跟踪正在工作的协程数量。当所有协程执行完毕（wg.done()），且已经处理完任务管道中的所有结果后，主协程才会关闭管道。这样就可以确保协程不会提前退出。