Golang 协程数量控制

写了个简单的端口开放探测,垃圾扫描器

大概流程是开始分两个子协程,一个把1-1024端口插入job管道,接着第二个子协程监听result管道,有内容时取出并打印。

接着往下是通过sync.WaitGroup创建指定数量的worker协程,worker的工作内容就是取job -> 探测端口 -> 结果插入result管道。

package main

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

type Job struct {
	host string
	port int
}

type Result struct {
	job    Job
	status bool
}

var jobs = make(chan Job, 1024)
var results = make(chan Result, 1024)

func worker(wg *sync.WaitGroup) {
	for job := range jobs {
		_, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", job.host, job.port), time.Second*3)
		if err != nil {
			results <- Result{job, false}
		} else {
			results <- Result{job, true}
		}
	}
	wg.Done()
}

func main() {
	sTime := time.Now()

	host := os.Args[1]
	// 创建任务
	go func() {
		for i := 1; i < 1024; i++ {
			jobs <- Job{host, i}
		}
		close(jobs)
	}()

	// 打印结果
	go func() {
		for result := range results {
			if result.status {
				fmt.Println(result.job, "open")
			}
		}
	}()

	// 创建一定数量的协程
	wg := sync.WaitGroup{}
	for i := 1; i < 200; i++ {
		wg.Add(1)
		go worker(&wg)
	}
	wg.Wait()
	//close(results)

	fmt.Println(time.Since(sTime))
}

发表评论

电子邮件地址不会被公开。 必填项已用*标注